from __future__ import annotations

import asyncio
import base64
import json
from collections.abc import Awaitable, Callable
from pathlib import Path
from typing import Any

import jinja2
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response

from device.manager import DeviceManager
from host_agent.assignment import AssignmentExecutor
from host_agent.client import (
    HostAgentClient,
    HostAgentEnrollmentClient,
    HostAgentAPIError,
    HostTaskSubmissionUnknownError,
)
from host_agent.config import HostAgentConfig
from host_agent.devices import register_local_device, unregister_local_device
from host_agent.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityStore
from host_agent.local_account import LocalAccountStore
from host_agent.status import AgentStatusTracker
from host_agent.web.auth import (
    SessionManager,
    SessionState,
    attempt_login,
    change_password,
)
from storage.device_config import DeviceConfigStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline

SESSION_COOKIE_NAME = "host_console_session"
CSRF_HEADER_NAME = "X-CSRF-Token"
CSRF_FORM_FIELD = "csrf_token"
_LOOPBACK_BIND_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"})

TaskSubmissionCallable = Callable[..., Awaitable[str]]
AUTOMATIC_DEVICE_VALUE = "__automatic__"

_ENV = jinja2.Environment(
    loader=jinja2.FileSystemLoader(Path(__file__).parent / "templates"),
    autoescape=jinja2.select_autoescape(["html", "xml"]),
)


def _render(
    template_name: str,
    *,
    status_code: int = 200,
    **context: Any,
) -> HTMLResponse:
    html = _ENV.get_template(template_name).render(**context)
    return HTMLResponse(html, status_code=status_code)


def _device_display_status(device: Any, *, busy_device_id: str | None) -> str:
    """Connected-but-idle devices report "busy" at the DeviceManager layer
    (an active Appium/WDA session is required to reuse it), which is not the
    same as a task currently running on that device. Only the device actually
    bound to the current assignment should read as "busy" here.
    """
    if device.status == "busy" and device.id != busy_device_id:
        return "connected"
    return device.status


def _screenshot_data_uri(
    record: dict[str, Any],
    *,
    path_key: str = "screenshot_path",
) -> str | None:
    """Return a ``data:`` URI for one step screenshot, or ``None``."""
    screenshot_path = record.get(path_key)
    if not screenshot_path:
        return None
    path = Path(str(screenshot_path))
    if not path.exists():
        return None
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def _ocr_results(record: dict[str, Any]) -> list[dict[str, Any]]:
    raw_results = record.get("ocr_results")
    if not isinstance(raw_results, list):
        return []
    return [result for result in raw_results if isinstance(result, dict)]


def _ui_tree_nodes(record: dict[str, Any]) -> list[dict[str, Any]]:
    tool_call = record.get("tool_call")
    if not isinstance(tool_call, dict):
        return []
    if tool_call.get("action") not in {"get_ui_tree", "ui_tree"}:
        return []
    step_result = record.get("result")
    if not isinstance(step_result, dict):
        return []
    raw_nodes = step_result.get("result")
    if not isinstance(raw_nodes, list):
        return []
    return [node for node in raw_nodes if isinstance(node, dict)]


def _timeline_step_context(record: dict[str, Any]) -> dict[str, Any]:
    tool_call = record.get("tool_call")
    result = record.get("result")
    return {
        "index": record.get("index", ""),
        "timestamp": record.get("timestamp", ""),
        "prompt": record.get("prompt") or "",
        "tool_call": tool_call if isinstance(tool_call, dict) else {},
        "result": result if isinstance(result, dict) else {},
        "before_screenshot_src": _screenshot_data_uri(
            record, path_key="before_screenshot_path"
        ),
        "after_screenshot_src": _screenshot_data_uri(
            record, path_key="after_screenshot_path"
        )
        or _screenshot_data_uri(record),
        "ocr_results": _ocr_results(record),
        "ui_tree_nodes": _ui_tree_nodes(record),
    }


def _safe_submission_error(detail: str) -> str:
    """Return a safe, single-line error message for the operator.

    The Cloud response ``detail`` is treated as a static control-plane message;
    we strip surrounding whitespace and reject empty results so the operator
    never sees a blank error or, through Jinja autoescape, anything that could
    carry unrendered HTML.
    """
    cleaned = " ".join(str(detail).split()).strip()
    return cleaned or "Cloud rejected the task submission."


def _extract_task_id(response: Any) -> str | None:
    """Normalize the Host self-submission return value to a Cloud task ID.

    Tests and narrow protocol overrides may return a bare string while the
    production client returns a pydantic model. Accept either so the rest of
    the handler can rely on a single string ID.
    """
    candidate: Any = response
    if hasattr(candidate, "task_id"):
        candidate = getattr(candidate, "task_id")
    if not isinstance(candidate, str):
        return None
    cleaned = candidate.strip()
    return cleaned or None


def _dashboard_texts(*, snapshot: dict[str, Any]) -> dict[str, str]:
    """Pre-compute human-readable text strings for the dashboard template."""
    heartbeat = snapshot.get("last_heartbeat")
    assignment = snapshot.get("current_assignment")
    policy = snapshot.get("host_policy")
    progress = snapshot.get("progress")
    return {
        "heartbeat_text": (
            f"{'ok' if heartbeat['ok'] else 'failed'} at {heartbeat['at']} "
            f"({heartbeat['device_count']} devices)"
            if heartbeat
            else "never"
        ),
        "assignment_text": (
            f"{assignment['task_id']} on {assignment['device_id']} "
            f"(started {assignment['started_at']})"
            if assignment
            else "none"
        ),
        "progress_text": (
            f"step {progress['step_index']} \u2014 {progress['step_status']}: "
            f"{progress['summary']}"
            if progress
            else ""
        ),
        "policy_text": (
            f"revision {policy['revision']}; self-submission "
            f"{'enabled' if policy['self_submission_enabled'] else 'disabled'}; "
            f"max active tasks {policy['max_active_tasks'] or 'unlimited'}; "
            f"daily token budget {policy['daily_token_budget'] or 'unmetered'}"
            if policy
            else "no Cloud policy cached"
        ),
    }


def create_console_app(
    *,
    config: HostAgentConfig,
    manager: DeviceManager,
    config_store: DeviceConfigStore,
    local_account_store: LocalAccountStore,
    identity_store: HostIdentityStore,
    history_store: ConsoleHistoryStore,
    status_tracker: AgentStatusTracker,
    session_manager: SessionManager,
    enrollment_client: HostAgentEnrollmentClient | None,
    host_client: HostAgentClient | None = None,
    submit_self_task: TaskSubmissionCallable | None = None,
    metadata_store: TaskMetadataStore | None = None,
    timeline: Timeline | None = None,
    executor: AssignmentExecutor | None = None,
) -> FastAPI:
    app = FastAPI(title="Host Agent Console")
    cookie_secure = config.console_bind_host not in _LOOPBACK_BIND_HOSTS
    if submit_self_task is None and host_client is not None:
        submit_self_task = host_client.submit_self_task
    submission_available = submit_self_task is not None

    def _running_devices() -> list[dict[str, str]]:
        return [
            {
                "id": device.id,
                "label": device.name or device.id,
            }
            for device in manager.list_devices()
        ]

    def _session_token(request: Request) -> str | None:
        return request.cookies.get(SESSION_COOKIE_NAME)

    def require_session(request: Request) -> SessionState:
        token = _session_token(request)
        state = session_manager.validate(token) if token else None
        if state is None:
            raise HTTPException(status_code=303, headers={"Location": "/login"})
        return state

    async def require_csrf(
        request: Request,
        session: SessionState = Depends(require_session),
    ) -> SessionState:
        provided = request.headers.get(CSRF_HEADER_NAME)
        if not provided:
            form = await request.form()
            raw = form.get(CSRF_FORM_FIELD)
            provided = str(raw) if raw is not None else None
        token = _session_token(request)
        if (
            not token
            or not provided
            or not session_manager.validate_csrf(token, provided)
        ):
            raise HTTPException(status_code=403, detail="invalid CSRF token")
        return session

    @app.get("/login", response_class=HTMLResponse)
    async def login_page() -> HTMLResponse:
        account = await asyncio.to_thread(local_account_store.load)
        return _render("login.html", title="Login", session=None, account=account)

    @app.post("/login")
    async def login_submit(request: Request) -> Response:
        account = await asyncio.to_thread(local_account_store.load)
        if account is None:
            return _render("login.html", title="Login", session=None, account=None)
        form = await request.form()
        username = str(form.get("username", ""))
        password = str(form.get("password", ""))
        ok = await asyncio.to_thread(
            attempt_login, local_account_store, username=username, password=password
        )
        if not ok:
            return _render(
                "login.html",
                title="Login",
                session=None,
                account=account,
                error="Invalid username or password.",
            )
        session_token, _ = session_manager.create_session(username)
        response = RedirectResponse(url="/", status_code=303)
        response.set_cookie(
            key=SESSION_COOKIE_NAME,
            value=session_token,
            httponly=True,
            samesite="strict",
            secure=cookie_secure,
            path="/",
        )
        return response

    @app.post("/logout")
    async def logout(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        token = _session_token(request)
        if token:
            session_manager.invalidate(token)
        response = RedirectResponse(url="/login", status_code=303)
        response.delete_cookie(key=SESSION_COOKIE_NAME, path="/")
        return response

    @app.get("/", response_class=HTMLResponse)
    async def dashboard(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        identity = await asyncio.to_thread(identity_store.load)
        snapshot = status_tracker.snapshot()
        busy_device_id = (
            snapshot["current_assignment"]["device_id"]
            if snapshot.get("current_assignment")
            else None
        )
        devices = [
            {
                "id": d.id,
                "name": d.name,
                "driver_type": d.driver_type,
                "display_status": _device_display_status(
                    d, busy_device_id=busy_device_id
                ),
            }
            for d in manager.list_devices()
        ]
        texts = _dashboard_texts(snapshot=snapshot)
        return _render(
            "dashboard.html",
            title="Status",
            session=session,
            identity=identity,
            devices=devices,
            config=config,
            **texts,
        )

    @app.get("/api/status")
    async def api_status(
        session: SessionState = Depends(require_session),
    ) -> JSONResponse:
        snapshot = status_tracker.snapshot()
        if executor is not None and snapshot.get("progress") is None:
            live = executor.latest_progress()
            if live is not None:
                snapshot["progress"] = {
                    "step_index": live.step_index,
                    "step_status": live.step_status,
                    "summary": live.summary,
                    "updated_at": live.updated_at.isoformat(),
                }
        current_assignment = snapshot.get("current_assignment")
        busy_device_id = current_assignment["device_id"] if current_assignment else None
        devices = [
            {
                "id": device.id,
                "name": device.name,
                "driver_type": device.driver_type,
                "status": _device_display_status(device, busy_device_id=busy_device_id),
            }
            for device in manager.list_devices()
        ]
        return JSONResponse({"status": snapshot, "devices": devices})

    @app.get("/devices", response_class=HTMLResponse)
    async def devices_page(
        request: Request,
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        devices = await asyncio.to_thread(config_store.list)
        edit_id = request.query_params.get("edit")
        edit_record = (
            await asyncio.to_thread(config_store.get, edit_id) if edit_id else None
        )
        connection_info_json = (
            json.dumps(edit_record["connection_info"]) if edit_record else "{}"
        )
        return _render(
            "devices.html",
            title="Devices",
            session=session,
            devices=devices,
            csrf_token=session.csrf_token,
            edit_record=edit_record,
            connection_info_json=connection_info_json,
            error=None,
        )

    @app.post("/devices/save")
    async def devices_save(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        form = await request.form()
        device_id = str(form.get("device_id", "")).strip()
        driver_type = str(form.get("driver_type", "")).strip()
        name = str(form.get("name", "")).strip() or None
        connection_info_raw = str(form.get("connection_info", "") or "{}")

        error: str | None = None
        connection_info: dict[str, Any] = {}
        if not device_id or not driver_type:
            error = "Device ID and driver type are required."
        else:
            try:
                parsed = json.loads(connection_info_raw)
            except ValueError:
                error = "Connection info must be valid JSON."
            else:
                if not isinstance(parsed, dict):
                    error = "Connection info must be a JSON object."
                else:
                    connection_info = parsed

        if error is None:
            try:
                await asyncio.to_thread(
                    register_local_device,
                    config_store,
                    manager,
                    device_id=device_id,
                    driver_type=driver_type,
                    connection_info=connection_info,
                    name=name,
                    config=config,
                    enrollment_client=enrollment_client,
                )
            except ValueError as exc:
                error = str(exc)

        if error is not None:
            devices = await asyncio.to_thread(config_store.list)
            return _render(
                "devices.html",
                title="Devices",
                session=session,
                devices=devices,
                csrf_token=session.csrf_token,
                edit_record=None,
                connection_info_json="{}",
                error=error,
                status_code=400,
            )
        return RedirectResponse(url="/devices", status_code=303)

    @app.post("/devices/remove")
    async def devices_remove(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        form = await request.form()
        device_id = str(form.get("device_id", "")).strip()
        if device_id:
            await asyncio.to_thread(
                unregister_local_device, config_store, manager, device_id=device_id
            )
        return RedirectResponse(url="/devices", status_code=303)

    @app.get("/account", response_class=HTMLResponse)
    async def account_page(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        return _render(
            "account.html",
            title="Account",
            session=session,
            csrf_token=session.csrf_token,
            message=None,
            error=None,
        )

    @app.post("/account", response_class=HTMLResponse)
    async def account_submit(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> HTMLResponse:
        form = await request.form()
        current_password = str(form.get("current_password", ""))
        new_password = str(form.get("new_password", ""))
        confirm_password = str(form.get("confirm_password", ""))
        if not new_password or new_password != confirm_password:
            return _render(
                "account.html",
                title="Account",
                session=session,
                csrf_token=session.csrf_token,
                message=None,
                error="New password and confirmation must match.",
                status_code=400,
            )
        ok = await asyncio.to_thread(
            change_password,
            local_account_store,
            current_password=current_password,
            new_password=new_password,
        )
        if not ok:
            return _render(
                "account.html",
                title="Account",
                session=session,
                csrf_token=session.csrf_token,
                message=None,
                error="Current password is incorrect.",
                status_code=400,
            )
        return _render(
            "account.html",
            title="Account",
            session=session,
            csrf_token=session.csrf_token,
            message="Password updated.",
            error=None,
        )

    @app.get("/history", response_class=HTMLResponse)
    async def history_page(
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        entries = await asyncio.to_thread(history_store.list_recent)
        return _render(
            "history.html",
            title="History",
            session=session,
            entries=entries,
        )

    def _tasks_list_context(
        session: SessionState,
        *,
        goal_value: str = "",
        selected_device: str = AUTOMATIC_DEVICE_VALUE,
        submission_error: str | None = None,
        submission_notice: str | None = None,
        submission_unknown: bool = False,
        status_code: int = 200,
    ) -> dict[str, Any]:
        devices = _running_devices()
        return {
            "title": "Tasks",
            "session": session,
            "csrf_token": session.csrf_token,
            "tasks": metadata_store.list_tasks() if metadata_store is not None else [],
            "metadata_store_missing": metadata_store is None,
            "devices": devices,
            "automatic_device_value": AUTOMATIC_DEVICE_VALUE,
            "selected_device": selected_device
            if any(d["id"] == selected_device for d in devices)
            or selected_device == AUTOMATIC_DEVICE_VALUE
            else AUTOMATIC_DEVICE_VALUE,
            "submission_available": submission_available,
            "goal_value": goal_value,
            "submission_error": submission_error,
            "submission_notice": submission_notice,
            "submission_unknown": submission_unknown,
            "status_code": status_code,
        }

    @app.get("/tasks", response_class=HTMLResponse)
    async def tasks_page(
        request: Request,
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        if metadata_store is None:
            raise HTTPException(
                status_code=503, detail="task metadata store not configured"
            )
        notice: str | None = None
        unknown = False
        error: str | None = None
        if request.query_params.get("submitted") == "1":
            task_id = request.query_params.get("task_id", "")
            if task_id:
                notice = (
                    f"Task submitted. Cloud task ID: {task_id}. "
                    "It will appear below when this Host begins executing it."
                )
        elif request.query_params.get("outcome") == "unknown":
            unknown = True
        context = _tasks_list_context(
            session,
            submission_notice=notice,
            submission_unknown=unknown,
            submission_error=error,
        )
        return _render(
            "tasks_list.html",
            status_code=context["status_code"],
            **{k: v for k, v in context.items() if k != "status_code"},
        )

    @app.post("/tasks/submit", response_class=HTMLResponse)
    async def tasks_submit(
        request: Request,
        session: SessionState = Depends(require_csrf),
    ) -> Response:
        if metadata_store is None:
            raise HTTPException(
                status_code=503, detail="task metadata store not configured"
            )
        if submit_self_task is None:
            context = _tasks_list_context(
                session,
                submission_error=(
                    "Host submission client is not available yet. "
                    "Wait for Host enrollment to complete, then retry."
                ),
            )
            return _render(
                "tasks_list.html",
                status_code=503,
                **{k: v for k, v in context.items() if k != "status_code"},
            )
        form = await request.form()
        raw_goal = str(form.get("goal", ""))
        goal = raw_goal.strip()
        device_selection = str(form.get("device_id", AUTOMATIC_DEVICE_VALUE)).strip()
        explicit_device_id: str | None = None
        if device_selection and device_selection != AUTOMATIC_DEVICE_VALUE:
            snapshot_ids = {device.id for device in manager.list_devices()}
            if device_selection not in snapshot_ids:
                context = _tasks_list_context(
                    session,
                    goal_value=goal,
                    selected_device=device_selection,
                    submission_error=(
                        "Selected device is no longer registered. "
                        "Refresh and try again."
                    ),
                )
                return _render(
                    "tasks_list.html",
                    status_code=400,
                    **{k: v for k, v in context.items() if k != "status_code"},
                )
            explicit_device_id = device_selection

        if not goal:
            context = _tasks_list_context(
                session,
                goal_value=goal,
                selected_device=device_selection or AUTOMATIC_DEVICE_VALUE,
                submission_error="Goal cannot be empty.",
            )
            return _render(
                "tasks_list.html",
                status_code=400,
                **{k: v for k, v in context.items() if k != "status_code"},
            )

        try:
            response = await submit_self_task(goal=goal, device_id=explicit_device_id)
        except HostAgentAPIError as exc:
            context = _tasks_list_context(
                session,
                goal_value=goal,
                selected_device=device_selection,
                submission_error=_safe_submission_error(str(exc.detail)),
            )
            return _render(
                "tasks_list.html",
                status_code=502,
                **{k: v for k, v in context.items() if k != "status_code"},
            )
        except HostTaskSubmissionUnknownError:
            return RedirectResponse(
                url="/tasks?outcome=unknown",
                status_code=303,
            )

        task_id = _extract_task_id(response)
        if task_id is None:
            context = _tasks_list_context(
                session,
                goal_value=goal,
                selected_device=device_selection,
                submission_error=(
                    "Host submission client returned an unexpected response."
                ),
            )
            return _render(
                "tasks_list.html",
                status_code=502,
                **{k: v for k, v in context.items() if k != "status_code"},
            )

        try:
            await asyncio.to_thread(
                history_store.record_task_submission,
                task_id=task_id,
                device_id=explicit_device_id,
            )
        except Exception:
            pass

        params = f"submitted=1&task_id={task_id}"
        return RedirectResponse(url=f"/tasks?{params}", status_code=303)

    @app.get("/tasks/{task_id}", response_class=HTMLResponse)
    async def task_detail_page(
        task_id: str,
        session: SessionState = Depends(require_session),
    ) -> HTMLResponse:
        if metadata_store is None:
            raise HTTPException(
                status_code=503, detail="task metadata store not configured"
            )
        task = await asyncio.to_thread(metadata_store.get_task, task_id)
        if task is None:
            raise HTTPException(status_code=404, detail="task not found")
        timeline_records: list[dict[str, Any]] = []
        if timeline is not None:
            timeline_records = await asyncio.to_thread(timeline.read, task_id)
        task_rows = [
            (label, task[key])
            for key, label in (
                ("source_task_id", "Cloud task ID"),
                ("source_attempt", "Cloud attempt"),
                ("id", "Execution ID"),
                ("goal", "Goal"),
                ("device_id", "Device"),
                ("status", "Status"),
                ("created_at", "Created"),
                ("updated_at", "Updated"),
                ("completed_at", "Completed"),
                ("failure_reason", "Failure reason"),
            )
            if task.get(key) is not None
        ]
        timeline_steps = [_timeline_step_context(record) for record in timeline_records]
        return _render(
            "task_detail.html",
            title=f"Task {task_id}",
            session=session,
            task=task,
            task_rows=task_rows,
            timeline_steps=timeline_steps,
        )

    return app