from __future__ import annotations

import asyncio
import logging
from contextlib import suppress
from dataclasses import dataclass

import uvicorn

from cloud.internal_api.models import AssignmentModel
from device.manager import DeviceManager
from host_agent.assignment import AssignmentExecutor
from host_agent.client import HostAgentClient, HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig, load_host_agent_config
from host_agent.dependency_supervisor import DependencySupervisor
from host_agent.devices import register_local_device
from host_agent.enrollment import resolve_host_identity
from host_agent.execution import create_execution_factories
from host_agent.heartbeat import HeartbeatSynchronizer
from host_agent.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityStore
from host_agent.instance_lock import InstanceLock
from host_agent.lease import ActiveAssignmentRunner
from host_agent.local_account import LocalAccountStore
from host_agent.mcp_lock import McpBusyTracker
from host_agent.mcp_token import McpTokenStore
from host_agent.policy_cache import HostPolicyCacheStore
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
from host_agent.retention import prune_task_history
from host_agent.skill_sync import HostAgentSkillSync
from host_agent.status import AgentStatusTracker
from host_agent.web.app import create_console_app
from host_agent.web.auth import SessionManager
from host_agent.web.mcp import build_mcp_server
from storage.artifact_store import ArtifactStore
from storage.device_config import DeviceConfigStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline


@dataclass
class HostAgentApplication:
    client: HostAgentClient
    heartbeat: HeartbeatSynchronizer
    processor: AssignmentProcessor
    console_server: uvicorn.Server | None = None
    console_enrollment_client: HostAgentEnrollmentClient | None = None
    dependency_supervisor: DependencySupervisor | None = None
    instance_lock: InstanceLock | None = None
    skill_sync: HostAgentSkillSync | None = None

    def run(self) -> None:
        asyncio.run(self.run_async())

    async def run_async(self, stop: asyncio.Event | None = None) -> None:
        stop_requested = stop or asyncio.Event()
        supervisor_stop = asyncio.Event()
        supervisor_task: asyncio.Task[None] | None = None
        if self.dependency_supervisor is not None:
            # Bring up supervised dependencies (probe + spawn/adopt + readiness
            # wait) before the heartbeat loop's first connect_devices() pass,
            # so a supervised Appium is ready before any Driver.connect().
            await self.dependency_supervisor.start()
            supervisor_task = asyncio.create_task(
                self.dependency_supervisor.run(supervisor_stop)
            )
        heartbeat_stop = asyncio.Event()
        heartbeat_task = asyncio.create_task(self.heartbeat.run(heartbeat_stop))
        if self.skill_sync is not None:
            self.skill_sync.start()
        console_task = (
            asyncio.create_task(self.console_server.serve())
            if self.console_server is not None
            else None
        )
        active_processing: asyncio.Task[AssignmentProcessingResult] | None = None
        try:
            while not stop_requested.is_set():
                assignment = await self._claim_until_stopped(stop_requested)
                if assignment is None:
                    continue
                active_processing = asyncio.create_task(
                    self.processor.process(assignment)
                )
                stopped = asyncio.create_task(stop_requested.wait())
                done, _ = await asyncio.wait(
                    {active_processing, stopped},
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if stopped in done:
                    self.processor.request_stop()
                else:
                    stopped.cancel()
                with suppress(asyncio.CancelledError):
                    await stopped
                await asyncio.shield(active_processing)
                active_processing = None
        finally:
            self.processor.request_stop()
            if active_processing is not None:
                with suppress(Exception):
                    await asyncio.shield(active_processing)
            heartbeat_stop.set()
            supervisor_stop.set()
            if self.console_server is not None:
                self.console_server.should_exit = True
            try:
                await asyncio.gather(heartbeat_task, return_exceptions=True)
                with suppress(Exception):
                    await self.heartbeat.sync_once()
                if supervisor_task is not None:
                    await asyncio.gather(supervisor_task, return_exceptions=True)
                    with suppress(Exception):
                        await self.dependency_supervisor.stop()
                if console_task is not None:
                    with suppress(asyncio.CancelledError):
                        await asyncio.gather(console_task, return_exceptions=True)
            finally:
                if self.console_enrollment_client is not None:
                    self.console_enrollment_client.close()
                if self.skill_sync is not None:
                    self.skill_sync.stop()
                await self.client.aclose()
                if self.instance_lock is not None:
                    self.instance_lock.release()

    async def _claim_until_stopped(
        self,
        stop: asyncio.Event,
    ) -> AssignmentModel | None:
        claim = asyncio.create_task(self.client.claim())
        stopped = asyncio.create_task(stop.wait())
        try:
            done, _ = await asyncio.wait(
                {claim, stopped},
                return_when=asyncio.FIRST_COMPLETED,
            )
        except asyncio.CancelledError:
            claim.cancel()
            stopped.cancel()
            await asyncio.gather(claim, stopped, return_exceptions=True)
            raise
        if stopped in done:
            claim.cancel()
            with suppress(asyncio.CancelledError):
                await claim
            return None
        stopped.cancel()
        with suppress(asyncio.CancelledError):
            await stopped
        return await claim


class _EmbeddedConsoleServer(uvicorn.Server):
    def install_signal_handlers(self) -> None:
        """Host Agent owns process-level signal handling."""


def create_application(
    *,
    config: HostAgentConfig | None = None,
    manager: DeviceManager | None = None,
    device_config_store: DeviceConfigStore | None = None,
    identity_store: HostIdentityStore | None = None,
    enrollment_client: HostAgentEnrollmentClient | None = None,
) -> HostAgentApplication:
    startup_config = config or load_host_agent_config()
    instance_lock = InstanceLock(startup_config.identity_path.parent)
    instance_lock.acquire()
    try:
        config_store = device_config_store or DeviceConfigStore()
        resolved_identity_store = identity_store or HostIdentityStore(
            startup_config.identity_path
        )
        owned_enrollment_client = enrollment_client is None
        bootstrap_client = enrollment_client or HostAgentEnrollmentClient(
            startup_config
        )
        try:
            resolved_config = resolve_host_identity(
                startup_config,
                identity_store=resolved_identity_store,
                client=bootstrap_client,
            )
            bootstrap_client.config = resolved_config
            resolved_manager = manager or _configured_device_manager(
                config_store,
                config=resolved_config,
                enrollment_client=bootstrap_client,
            )
        finally:
            if owned_enrollment_client:
                bootstrap_client.close()
        client = HostAgentClient(resolved_config)

        history_store = ConsoleHistoryStore(
            resolved_config.identity_path.parent / "host_console_history.sqlite3",
            limit=resolved_config.console_history_limit,
        )
        status_tracker = AgentStatusTracker()
        console_enrollment_client: HostAgentEnrollmentClient | None = None
        if resolved_config.enrollment_managed:
            console_enrollment_client = HostAgentEnrollmentClient(resolved_config)
        metadata_store = TaskMetadataStore(
            db_path=resolved_config.task_progress_db_path
        )
        timeline = Timeline(ArtifactStore(root=resolved_config.task_artifact_dir))
        mcp_token_path = resolved_config.identity_path.parent / "host_mcp_token.json"
        mcp_token_existed = mcp_token_path.exists()
        mcp_token_store = McpTokenStore(mcp_token_path)
        mcp_token_store.load_or_create()
        if not mcp_token_existed:
            logging.getLogger(__name__).info(
                "MCP token generated at %s", mcp_token_path
            )
        mcp_busy_tracker = McpBusyTracker(ttl_seconds=20.0)
        executor = AssignmentExecutor(
            create_execution_factories(
                resolved_manager,
                metadata_store=metadata_store,
                timeline=timeline,
                host_agent_config=resolved_config,
            ),
            mcp_busy_tracker=mcp_busy_tracker,
        )
        mcp_server = build_mcp_server(
            manager=resolved_manager,
            mcp_busy_tracker=mcp_busy_tracker,
            status_tracker=status_tracker,
        )
        console_app = create_console_app(
            config=resolved_config,
            manager=resolved_manager,
            config_store=config_store,
            local_account_store=LocalAccountStore(resolved_config.local_account_path),
            identity_store=resolved_identity_store,
            history_store=history_store,
            status_tracker=status_tracker,
            session_manager=SessionManager(
                ttl_seconds=resolved_config.console_session_ttl_seconds
            ),
            enrollment_client=console_enrollment_client,
            host_client=client,
            metadata_store=metadata_store,
            timeline=timeline,
            executor=executor,
            mcp_server=mcp_server,
            mcp_token_store=mcp_token_store,
            mcp_busy_tracker=mcp_busy_tracker,
        )
        console_server = _EmbeddedConsoleServer(
            uvicorn.Config(
                console_app,
                host=resolved_config.console_bind_host,
                port=resolved_config.console_port,
                log_level="warning",
            )
        )

        heartbeat = HeartbeatSynchronizer(
            resolved_manager,
            client,
            resolved_config,
            status_tracker=status_tracker,
            mcp_busy_tracker=mcp_busy_tracker,
            on_sync=lambda device_count: history_store.record_heartbeat(
                device_count=device_count
            ),
            policy_cache=HostPolicyCacheStore(
                resolved_config.identity_path.parent / "host_governance_policy.json"
            ),
            on_policy_sync=lambda revision: history_store.record_policy_sync(
                revision=revision
            ),
        )
        active_runner = ActiveAssignmentRunner(client, executor)
        processor = AssignmentProcessor(
            client,
            active_runner,
            status_tracker=status_tracker,
            on_result=lambda assignment, result: _on_assignment_finished(
                history_store,
                metadata_store,
                timeline,
                resolved_config,
                assignment,
                result,
            ),
        )
        dependency_supervisor: DependencySupervisor | None = None
        if resolved_config.dependency_supervisor_enabled:
            dependency_supervisor = DependencySupervisor.from_host_agent_config(
                resolved_config
            )
        skill_sync: HostAgentSkillSync | None = None
        if resolved_config.host_id and resolved_config.token:
            skill_sync = HostAgentSkillSync(resolved_config)
        return HostAgentApplication(
            client=client,
            heartbeat=heartbeat,
            processor=processor,
            console_server=console_server,
            console_enrollment_client=console_enrollment_client,
            dependency_supervisor=dependency_supervisor,
            instance_lock=instance_lock,
            skill_sync=skill_sync,
        )
    except BaseException:
        instance_lock.release()
        raise


def _on_assignment_finished(
    history_store: ConsoleHistoryStore,
    metadata_store: TaskMetadataStore,
    timeline: Timeline,
    config: HostAgentConfig,
    assignment: AssignmentModel,
    result: AssignmentProcessingResult,
) -> None:
    _record_assignment_history(history_store, assignment, result)
    _prune_task_history(metadata_store, timeline, config)


def _record_assignment_history(
    history_store: ConsoleHistoryStore,
    assignment: AssignmentModel,
    result: AssignmentProcessingResult,
) -> None:
    status = "done" if result.execution.status == "done" else "failed"
    history_store.record_assignment(
        task_id=assignment.task_id,
        attempt=assignment.attempt,
        status=status,
        failure_reason=result.execution.failure_reason if status == "failed" else None,
        device_id=assignment.device_id,
    )


def _prune_task_history(
    metadata_store: TaskMetadataStore,
    timeline: Timeline,
    config: HostAgentConfig,
) -> None:
    try:
        prune_task_history(metadata_store, timeline, config=config)
    except Exception:
        pass


def _configured_device_manager(
    config_store: DeviceConfigStore,
    *,
    config: HostAgentConfig,
    enrollment_client: HostAgentEnrollmentClient,
) -> DeviceManager:
    manager = DeviceManager()
    for device in config_store.list():
        register_local_device(
            config_store,
            manager,
            device_id=device["device_id"],
            driver_type=device["driver_type"],
            connection_info=device["connection_info"],
            name=device["name"],
            config=config,
            enrollment_client=enrollment_client,
        )
    return manager