#!/usr/bin/env python3
"""Hand project-scoped queued investigations to the Agent.

Credentials are deployment environment values, never task/session payloads.
Run with --once for an operator-triggered tick or without it for polling.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import time
from dataclasses import dataclass, replace
from datetime import datetime
from pathlib import Path
from uuid import UUID, uuid4

import httpx
from dotenv import load_dotenv


logger = logging.getLogger("operational-worker")


@dataclass(frozen=True)
class WorkerSettings:
    server_url: str
    agent_url: str
    project_id: str
    interval_seconds: float = 30
    max_concurrency: int = 4
    bearer: str = ""
    token_url: str = ""
    client_id: str = ""
    client_secret: str = ""
    token_refresh_seconds: float = 300
    task_created_from: str | None = None
    project_ids: tuple[str, ...] = ()

    def project_scopes(self) -> tuple["WorkerSettings", ...]:
        """Keep credentials, sessions and cursors separate for each explicit project."""
        return tuple(replace(self, project_id=project_id, project_ids=())
                     for project_id in self.project_ids or (self.project_id,))

    @classmethod
    def from_environment(cls) -> "WorkerSettings":
        load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False)
        configured_projects = os.environ.get("OPERATIONAL_WORKER_PROJECT_IDS", "").strip()
        project_ids = tuple(dict.fromkeys(
            str(UUID(value.strip())) for value in configured_projects.split(",")
        )) if configured_projects else ()
        project_id = project_ids[0] if project_ids else str(UUID(os.environ["OPERATIONAL_WORKER_PROJECT_ID"]))
        interval = float(os.environ.get("OPERATIONAL_WORKER_RECONCILE_SECONDS", "30"))
        if interval < 1:
            raise ValueError("worker interval must be at least one second")
        token_url = os.environ.get("OPERATIONAL_WORKER_TOKEN_URL", "")
        if os.environ.get("ENVIRONMENT") == "development":
            token_url = os.environ.get("OPERATIONAL_WORKER_DEV_TOKEN_URL") or token_url
        settings = cls(
            server_url=os.environ.get("OPERATIONAL_WORKER_SERVER_URL", "http://127.0.0.1:8000"),
            agent_url=os.environ.get("OPERATIONAL_WORKER_AGENT_URL", "http://127.0.0.1:8787"),
            project_id=project_id,
            interval_seconds=interval,
            max_concurrency=int(os.environ.get("OPERATIONAL_WORKER_MAX_CONCURRENCY", "4")),
            bearer=os.environ.get("OPERATIONAL_WORKER_BEARER", ""),
            token_url=token_url,
            client_id=os.environ.get("OPERATIONAL_WORKER_CLIENT_ID", ""),
            client_secret=os.environ.get("OPERATIONAL_WORKER_CLIENT_SECRET", ""),
            token_refresh_seconds=float(os.environ.get("OPERATIONAL_WORKER_TOKEN_REFRESH_SECONDS", "300")),
            task_created_from=os.environ.get("OPERATIONAL_WORKER_TASK_CREATED_FROM") or None,
            project_ids=project_ids,
        )
        if settings.task_created_from:
            boundary = datetime.fromisoformat(settings.task_created_from.replace("Z", "+00:00"))
            if boundary.tzinfo is None:
                raise ValueError("worker task creation boundary must include a timezone")
        if not 30 <= settings.token_refresh_seconds <= 600:
            raise ValueError("service token refresh interval must be between 30 and 600 seconds")
        if not 1 <= settings.max_concurrency <= 32:
            raise ValueError("worker concurrency must be between 1 and 32")
        if not settings.bearer and not all((settings.token_url, settings.client_id, settings.client_secret)):
            raise ValueError("configure a product service credential for the operational worker")
        return settings


class CredentialProvider:
    def __init__(self, settings: WorkerSettings):
        self.settings = settings
        self._token = ""
        self._expires = 0.0

    async def headers(self, client: httpx.AsyncClient) -> dict[str, str]:
        if self.settings.bearer:
            token = self.settings.bearer.removeprefix("Bearer ")
        else:
            if time.monotonic() >= self._expires:
                response = await client.post(self.settings.token_url, data={
                    "grant_type": "client_credentials", "client_id": self.settings.client_id,
                    "client_secret": self.settings.client_secret,
                })
                response.raise_for_status()
                payload = response.json()
                self._token = payload["access_token"]
                if not isinstance(self._token, str) or not self._token:
                    raise ValueError("identity provider returned no access token")
                # Product auth also restricts token age, independently of the
                # provider's exp. Do not cache a long-lived OAuth token for hours.
                lifetime = min(float(payload.get("expires_in", 60)), self.settings.token_refresh_seconds)
                self._expires = time.monotonic() + max(1, lifetime - 15)
            token = self._token
        return {"Authorization": f"Bearer {token}", "X-Project-Id": self.settings.project_id}


class OperationalWorker:
    def __init__(self, settings: WorkerSettings, client: httpx.AsyncClient):
        self.settings = settings
        self.client = client
        self.credentials = CredentialProvider(settings)
        self._agent_sessions: dict[str, str] = {}
        self._message_cursor: str | None = None

    async def tick(self) -> int:
        headers = await self.credentials.headers(self.client)
        server = self.settings.server_url.rstrip("/")
        scope = {"task_type": "event_investigation"}
        if self.settings.task_created_from:
            scope["task_created_from"] = self.settings.task_created_from
        recovery = await self.client.post(
            f"{server}/api/v1/operational-tasks/leases/recover", headers=headers,
            json={"limit": 100, **scope},
        )
        recovery.raise_for_status()
        active_run_ids: set[str] = set()
        for status_value in ("claimed", "running"):
            response = await self.client.get(
                f"{server}/api/v1/operational-task-runs", headers=headers,
                params={"status": status_value, "limit": self.settings.max_concurrency, **scope},
            )
            response.raise_for_status()
            for item in response.json()["items"]:
                run_id = str(item["run_id"])
                active_run_ids.add(run_id)
                session_id = item.get("agent_session_id")
                if isinstance(session_id, str) and session_id:
                    self._agent_sessions[run_id] = session_id
        for run_id, session_id in list(self._agent_sessions.items()):
            if run_id not in active_run_ids:
                self._agent_sessions.pop(run_id, None)
                continue
            # Refresh the Agent vault with the worker's existing OAuth provider.
            # Background runs never depend on an interactive browser refresh.
            renewed = await self.client.post(
                f"{self.settings.agent_url.rstrip('/')}/api/v1/session/{session_id}/credentials",
                headers=headers, json={},
            )
            if renewed.status_code not in (202, 404, 409):
                renewed.raise_for_status()
        available = max(0, self.settings.max_concurrency - len(active_run_ids))
        if available == 0:
            await self._reconcile_messages(headers)
            return 0
        response = await self.client.get(
            f"{server}/api/v1/operational-task-runs", headers=headers,
            params={"status": "queued", "limit": available, **scope},
        )
        response.raise_for_status()
        queued = response.json()["items"]

        async def dispatch(run: dict) -> bool:
            if run.get("status") != "queued":
                return False
            task_id = str(UUID(run["task_id"]))
            run_id = str(UUID(run["run_id"]))
            request_id = str(UUID(run["execution_request_id"]))
            try:
                result = await self.client.post(
                    f"{self.settings.agent_url.rstrip('/')}/api/v1/task-executions",
                    headers=headers,
                    json={
                        "task_id": task_id,
                        "run_id": run_id,
                        "execution_request_id": request_id,
                    },
                )
                if result.is_success:
                    payload = result.json()
                    if payload.get("task_id") != task_id or payload.get("run_id") != run_id:
                        raise ValueError("Agent returned a different task run")
                    session_id = payload.get("session_id")
                    if isinstance(session_id, str) and session_id:
                        self._agent_sessions[run_id] = session_id
                    return True
                await self._record_dispatch_failure(
                    headers, run, retryable=result.status_code >= 500,
                    code=f"AGENT_HTTP_{result.status_code}",
                )
                return False
            except (httpx.HTTPError, ValueError) as exc:
                await self._record_dispatch_failure(
                    headers, run, retryable=True,
                    code="AGENT_UNAVAILABLE", detail=type(exc).__name__,
                )
                return False

        results = await asyncio.gather(*(dispatch(run) for run in queued))
        await self._reconcile_messages(headers)
        return sum(results)

    async def _reconcile_messages(self, headers: dict[str, str]) -> None:
        # Reconcile already bound runs even when task dispatch is at capacity.
        # Agent and Server share immutable message/native identities for retries.
        semaphore = asyncio.Semaphore(self.settings.max_concurrency)

        async def dispatch(message: dict) -> None:
            async with semaphore:
                message_id = str(UUID(message["message_id"]))
                try:
                    delivered = await self.client.post(
                        f"{self.settings.agent_url.rstrip('/')}/api/v1/task-messages/deliver",
                        headers=headers, json={"message_id": message_id},
                    )
                    if delivered.status_code in {400, 403, 404, 410, 422}:
                        failed = await self.client.post(
                            f"{self.settings.server_url.rstrip('/')}/api/v1/agent-messages/{message_id}/delivery",
                            headers=headers, json={
                                "status": "failed",
                                "failure_detail": f"Agent rejected message delivery (HTTP {delivered.status_code})",
                            },
                        )
                        if failed.status_code != 409:
                            failed.raise_for_status()
                    elif not delivered.is_success:
                        delivered.raise_for_status()
                except httpx.HTTPError:
                    # Avoid message content and credentials in logs. Durable
                    # nonterminal messages are available at the next tick.
                    logger.warning("Agent message delivery remains pending: %s", message_id)

        # Bounded pages and a cursor retained between ticks prevent idle chats
        # from occupying the first page forever. Completing a sweep starts the
        # next tick from the beginning so older transient failures are retried.
        for _ in range(4):
            params = {"limit": 50}
            if self._message_cursor:
                params["cursor"] = self._message_cursor
            response = await self.client.get(
                f"{self.settings.server_url.rstrip('/')}/api/v1/agent-messages/pending",
                headers=headers, params=params,
            )
            response.raise_for_status()
            payload = response.json()
            await asyncio.gather(*(dispatch(message) for message in payload["items"]))
            self._message_cursor = payload.get("next_cursor")
            if not self._message_cursor:
                break

    async def _record_dispatch_failure(
        self, headers: dict[str, str], run: dict, *, retryable: bool,
        code: str, detail: str = "Agent rejected task run dispatch",
    ) -> None:
        response = await self.client.post(
            f"{self.settings.server_url.rstrip('/')}/api/v1/operational-tasks/"
            f"{run['task_id']}/runs/{run['run_id']}/dispatch-failure",
            headers=headers,
            json={
                "request_id": str(uuid4()),
                "expected_revision": run["revision"],
                "failure_code": code,
                "failure_detail": detail,
                "retryable": retryable,
            },
        )
        response.raise_for_status()

    async def listen(self, wake: asyncio.Event) -> None:
        cursor = 0
        while True:
            try:
                headers = await self.credentials.headers(self.client)
                async with self.client.stream(
                    "GET",
                    f"{self.settings.server_url.rstrip('/')}/api/v1/workbench/events",
                    headers={**headers, "Accept": "text/event-stream"},
                    params={"after": cursor},
                    timeout=None,
                ) as response:
                    response.raise_for_status()
                    event_name = ""
                    reset_required = False
                    async for line in response.aiter_lines():
                        if line.startswith("id:"):
                            cursor = int(line[3:].strip())
                        elif line.startswith("event:"):
                            event_name = line[6:].strip()
                        elif line.startswith("data:") and event_name == "reset_required":
                            snapshot = await self.client.get(
                                f"{self.settings.server_url.rstrip('/')}/api/v1/workbench/sync-snapshot",
                                headers=headers,
                            )
                            snapshot.raise_for_status()
                            cursor = int(snapshot.json()["event_cursor"])
                            reset_required = True
                            break
                        elif line.startswith("data:") and event_name == "workbench_event":
                            payload = json.loads(line[5:].strip())
                            if payload.get("topic") == "operational.run.queued":
                                wake.set()
                        elif line == "":
                            event_name = ""
                    if not reset_required:
                        await asyncio.sleep(1)
            except (httpx.HTTPError, ValueError, KeyError):
                logger.warning("Workbench event stream disconnected; reconciliation remains active")
                await asyncio.sleep(5)


async def run_project(settings: WorkerSettings, client: httpx.AsyncClient, once: bool) -> None:
    logger.info(
        "Investigation worker scope: project=%s task_created_from=%s concurrency=%d",
        settings.project_id, settings.task_created_from or "all", settings.max_concurrency,
    )
    worker = OperationalWorker(settings, client)
    wake = asyncio.Event()
    listener = None if once else asyncio.create_task(worker.listen(wake))
    try:
        while True:
            try:
                accepted = await worker.tick()
                logger.info("Task dispatch tick completed: project=%s accepted=%d", settings.project_id, accepted)
            except (httpx.HTTPError, ValueError, KeyError):
                logger.error("Task dispatch tick failed: project=%s; persisted tasks remain available for retry", settings.project_id)
                if once:
                    raise
            if once:
                return
            try:
                await asyncio.wait_for(wake.wait(), timeout=settings.interval_seconds)
            except TimeoutError:
                pass
            wake.clear()
    finally:
        if listener is not None:
            listener.cancel()
            await asyncio.gather(listener, return_exceptions=True)


async def run(once: bool) -> None:
    settings = WorkerSettings.from_environment()
    async with httpx.AsyncClient(timeout=30, follow_redirects=False, trust_env=False) as client:
        results = await asyncio.gather(
            *(run_project(scope, client, once) for scope in settings.project_scopes()),
            return_exceptions=once,
        )
        if any(isinstance(result, Exception) for result in results):
            raise SystemExit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO)
    try:
        asyncio.run(run(args.once))
    except KeyboardInterrupt:
        pass