#!/usr/bin/env python3
"""HTTP API for AIOS external-agent online task distribution.

The service is intended to bind to loopback and be published through the
existing TLS reverse proxy. Control operations use a server-managed control
Bearer token. Agents receive individual Bearer tokens at registration; only
token hashes are persisted.
"""
from __future__ import annotations

from dataclasses import dataclass
import hmac
import json
import os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import re
from typing import Any, Mapping
from urllib.parse import urlparse

from aios_agent_task_protocol import TaskProtocolError, TaskStore


MAX_BODY = 1024 * 1024
TASK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{2,179}$")


@dataclass(frozen=True, slots=True)
class APIConfig:
    db_path: Path
    control_token_path: Path
    host: str = "127.0.0.1"
    port: int = 18940
    token_pepper: str = ""

    @classmethod
    def from_env(cls) -> "APIConfig":
        return cls(
            db_path=Path(os.getenv(
                "AIOS_AGENT_TASK_DB",
                "/var/lib/aios-agent-task/state.sqlite3",
            )),
            control_token_path=Path(os.getenv(
                "AIOS_AGENT_TASK_CONTROL_TOKEN_FILE",
                "/run/credentials/aios-agent-task-api.service/control_token",
            )),
            host=os.getenv("AIOS_AGENT_TASK_HOST", "127.0.0.1"),
            port=int(os.getenv("AIOS_AGENT_TASK_PORT", "18940")),
            token_pepper=os.getenv("AIOS_AGENT_TASK_TOKEN_PEPPER", ""),
        )


class AgentTaskApplication:
    def __init__(self, config: APIConfig) -> None:
        self.config = config
        self.store = TaskStore(
            config.db_path,
            token_pepper=config.token_pepper,
        )

    def _control_token(self) -> str:
        path = self.config.control_token_path
        if path.is_symlink() or not path.is_file():
            raise TaskProtocolError("CONTROL_CREDENTIAL_UNAVAILABLE", 503)
        value = path.read_text(encoding="utf-8").strip()
        if len(value) < 32:
            raise TaskProtocolError("CONTROL_CREDENTIAL_INVALID", 503)
        return value

    @staticmethod
    def _bearer(headers: Mapping[str, str]) -> str:
        value = str(headers.get("Authorization") or "")
        if not value.startswith("Bearer "):
            return ""
        return value[7:].strip()

    def authenticate_control(self, headers: Mapping[str, str]) -> None:
        supplied = self._bearer(headers)
        expected = self._control_token()
        if not supplied or not hmac.compare_digest(supplied, expected):
            raise TaskProtocolError("UNAUTHORIZED", 401)

    def authenticate_agent(self, headers: Mapping[str, str]) -> dict[str, Any]:
        return self.store.authenticate_agent(self._bearer(headers))

    def principal(self, headers: Mapping[str, str]) -> tuple[str, dict[str, Any] | None]:
        supplied = self._bearer(headers)
        if not supplied:
            raise TaskProtocolError("UNAUTHORIZED", 401)
        try:
            expected = self._control_token()
        except TaskProtocolError:
            expected = ""
        if expected and hmac.compare_digest(supplied, expected):
            return "control", None
        return "agent", self.store.authenticate_agent(supplied)

    def health(self) -> dict[str, Any]:
        return {
            "status": "ok",
            "schema": "aios.agent-task.v1",
            "task_states": [
                "QUEUED", "CLAIMED", "RUNNING", "VERIFYING",
                "SUCCEEDED", "FAILED", "EXPIRED", "CANCELLED",
            ],
            "lease_recovery": True,
            "atomic_claim": True,
            "idempotent_result_submission": True,
            "server_side_verification_required": True,
            "plaintext_credentials_persisted": False,
        }

    def register_agent(self, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        self.authenticate_control(headers)
        return self.store.register_agent(
            agent_id=str(body.get("agent_id") or ""),
            agent_type=str(body.get("agent_type") or ""),
            agent_version=str(body.get("agent_version") or ""),
            capabilities=body.get("capabilities") or [],
            metadata=body.get("metadata") if isinstance(body.get("metadata"), Mapping) else {},
        )

    def create_task(self, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        self.authenticate_control(headers)
        return self.store.create_task(
            task_id=str(body.get("task_id") or "") or None,
            goal=str(body.get("goal") or ""),
            instructions=str(body.get("instructions") or ""),
            success_criteria=body.get("success_criteria"),
            required_capabilities=body.get("required_capabilities") or [],
            constraints=body.get("constraints") if isinstance(body.get("constraints"), Mapping) else {},
            metadata=body.get("metadata") if isinstance(body.get("metadata"), Mapping) else {},
            deadline=str(body.get("deadline") or "") or None,
            priority=body.get("priority", 0),
            max_attempts=body.get("max_attempts", 3),
            idempotency_key=str(body.get("idempotency_key") or ""),
        )

    def poll(self, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        agent = self.authenticate_agent(headers)
        tasks = self.store.poll_tasks(
            agent_id=agent["agent_id"],
            limit=int(body.get("limit", 20)),
        )
        return {"agent_id": agent["agent_id"], "tasks": tasks}

    def claim(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        agent = self.authenticate_agent(headers)
        return self.store.claim_task(
            task_id=task_id,
            agent_id=agent["agent_id"],
            lease_seconds=int(body.get("lease_seconds", 300)),
        )

    def heartbeat(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        agent = self.authenticate_agent(headers)
        return self.store.heartbeat(
            task_id=task_id,
            agent_id=agent["agent_id"],
            lease_token=str(body.get("lease_token") or ""),
            progress=body.get("progress") if isinstance(body.get("progress"), Mapping) else {},
            lease_seconds=int(body.get("lease_seconds", 300)),
        )

    def result(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        agent = self.authenticate_agent(headers)
        result = body.get("result")
        if not isinstance(result, Mapping):
            raise TaskProtocolError("RESULT_INVALID", 400)
        return self.store.submit_result(
            task_id=task_id,
            agent_id=agent["agent_id"],
            lease_token=str(body.get("lease_token") or ""),
            idempotency_key=str(body.get("idempotency_key") or ""),
            result=result,
        )

    def verify(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        self.authenticate_control(headers)
        return self.store.verify_task(
            task_id=task_id,
            decision=str(body.get("decision") or ""),
            verification=body.get("verification") if isinstance(body.get("verification"), Mapping) else {},
        )

    def cancel(self, task_id: str, headers: Mapping[str, str], body: Mapping[str, Any]) -> dict[str, Any]:
        self.authenticate_control(headers)
        return self.store.cancel_task(task_id=task_id, reason=str(body.get("reason") or ""))

    def get_task(self, task_id: str, headers: Mapping[str, str]) -> dict[str, Any]:
        kind, agent = self.principal(headers)
        task = self.store.get_task(task_id)
        if kind == "agent" and task.get("claimed_by") != agent["agent_id"]:
            raise TaskProtocolError("FORBIDDEN", 403)
        return task

    def events(self, task_id: str, headers: Mapping[str, str]) -> dict[str, Any]:
        self.authenticate_control(headers)
        return {"task_id": task_id, "events": self.store.get_events(task_id)}


class Handler(BaseHTTPRequestHandler):
    app: AgentTaskApplication | None = None
    server_version = "AIOSAgentTaskAPI/1"

    def log_message(self, _format: str, *_args: Any) -> None:
        return

    def _headers(self) -> dict[str, str]:
        return {key: value for key, value in self.headers.items()}

    def _json(self, value: Any, status: int = 200) -> None:
        payload = json.dumps(value, ensure_ascii=False, sort_keys=True).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def _body(self) -> dict[str, Any]:
        try:
            size = int(self.headers.get("Content-Length", "0"))
        except ValueError as exc:
            raise TaskProtocolError("CONTENT_LENGTH_INVALID", 400) from exc
        if size < 0 or size > MAX_BODY:
            raise TaskProtocolError("REQUEST_TOO_LARGE", 413)
        raw = self.rfile.read(size) if size else b"{}"
        try:
            value = json.loads(raw.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise TaskProtocolError("JSON_INVALID", 400) from exc
        if not isinstance(value, dict):
            raise TaskProtocolError("JSON_OBJECT_REQUIRED", 400)
        return value

    @staticmethod
    def _task_route(path: str) -> tuple[str, str | None]:
        match = re.fullmatch(
            r"/v1/tasks/([A-Za-z0-9][A-Za-z0-9._:-]{2,179})(?:/(claim|heartbeat|result|verify|cancel|events))?",
            path,
        )
        if match:
            return match.group(1), match.group(2)
        return "", None

    def _run(self, callback) -> None:
        try:
            callback()
        except TaskProtocolError as exc:
            self._json({"error": exc.code}, exc.http_status)
        except Exception as exc:
            self._json({"error": "INTERNAL_ERROR", "detail": type(exc).__name__}, 500)

    def do_GET(self) -> None:
        def action() -> None:
            if self.app is None:
                raise TaskProtocolError("SERVICE_UNAVAILABLE", 503)
            path = urlparse(self.path).path.rstrip("/") or "/"
            if path == "/health":
                self._json(self.app.health())
                return
            task_id, suffix = self._task_route(path)
            if task_id and suffix is None:
                self._json(self.app.get_task(task_id, self._headers()))
                return
            if task_id and suffix == "events":
                self._json(self.app.events(task_id, self._headers()))
                return
            self._json({"error": "NOT_FOUND"}, 404)
        self._run(action)

    def do_POST(self) -> None:
        def action() -> None:
            if self.app is None:
                raise TaskProtocolError("SERVICE_UNAVAILABLE", 503)
            path = urlparse(self.path).path.rstrip("/") or "/"
            body = self._body()
            if path == "/v1/agents/register":
                self._json(self.app.register_agent(self._headers(), body), 201)
                return
            if path == "/v1/tasks":
                self._json(self.app.create_task(self._headers(), body), 201)
                return
            if path == "/v1/tasks/poll":
                self._json(self.app.poll(self._headers(), body))
                return
            task_id, suffix = self._task_route(path)
            if not task_id:
                self._json({"error": "NOT_FOUND"}, 404)
                return
            if suffix == "claim":
                self._json(self.app.claim(task_id, self._headers(), body))
            elif suffix == "heartbeat":
                self._json(self.app.heartbeat(task_id, self._headers(), body))
            elif suffix == "result":
                self._json(self.app.result(task_id, self._headers(), body))
            elif suffix == "verify":
                self._json(self.app.verify(task_id, self._headers(), body))
            elif suffix == "cancel":
                self._json(self.app.cancel(task_id, self._headers(), body))
            else:
                self._json({"error": "NOT_FOUND"}, 404)
        self._run(action)


def serve(config: APIConfig | None = None) -> None:
    config = config or APIConfig.from_env()
    application = AgentTaskApplication(config)
    Handler.app = application
    ThreadingHTTPServer((config.host, config.port), Handler).serve_forever()


if __name__ == "__main__":
    serve()