from __future__ import annotations

import time
import uuid
from typing import Any, cast
from urllib.parse import quote

import httpx

from grantflow.api.routes.agents import _agent_tool_manifest_payload

DEFAULT_BASE_URL = "http://127.0.0.1:8000"
GRANTFLOW_TOOL_PREFIX = "grantflow_"
TERMINAL_STATUSES = {"done", "error", "pending_hitl"}


class GrantFlowToolError(RuntimeError):
    def __init__(self, message: str, *, detail: dict[str, Any] | None = None) -> None:
        super().__init__(message)
        self.detail = detail or {"message": message}


def _prefixed_name(name: str) -> str:
    return f"{GRANTFLOW_TOOL_PREFIX}{name}"


def _unprefixed_name(name: str) -> str:
    if name.startswith(GRANTFLOW_TOOL_PREFIX):
        return name.removeprefix(GRANTFLOW_TOOL_PREFIX)
    return name


def list_tools() -> list[dict[str, Any]]:
    manifest = _agent_tool_manifest_payload()
    manifest_tools = cast(list[dict[str, Any]], manifest["tools"])
    tools: list[dict[str, Any]] = []
    for tool in manifest_tools:
        input_schema = tool.get("input_schema", {})
        tools.append(
            {
                "name": _prefixed_name(str(tool["name"])),
                "description": tool.get("description", ""),
                "inputSchema": input_schema,
                "annotations": {
                    "grantflow_name": tool["name"],
                    "method": tool.get("method"),
                    "path": tool.get("path"),
                    "auth_required": tool.get("auth_required", False),
                    "idempotency_key_supported": tool.get("idempotency_key_supported", False),
                    "output_schema_ref": tool.get("output_schema_ref"),
                    "output_hints": tool.get("output_hints", []),
                },
            }
        )
    tools.append(
        {
            "name": "grantflow_run_sandbox_happy_path",
            "description": "Register a sandbox agent, run preflight, start generation, then read status, quality, and events.",
            "inputSchema": {
                "type": "object",
                "additionalProperties": False,
                "required": ["agent_name"],
                "properties": {
                    "agent_name": {"type": "string"},
                    "agent_type": {
                        "type": "string",
                        "enum": ["llm_agent", "workflow_agent", "integration", "human_operator"],
                        "default": "workflow_agent",
                    },
                    "purpose": {"type": "string"},
                    "poll_interval_seconds": {"type": "number", "default": 0.5},
                    "timeout_seconds": {"type": "number", "default": 20.0},
                },
            },
            "annotations": {"workflow": "sandbox_generate"},
        }
    )
    return tools


class GrantFlowToolClient:
    def __init__(
        self,
        *,
        base_url: str = DEFAULT_BASE_URL,
        api_key: str | None = None,
        bearer_token: str | None = None,
        client: httpx.Client | None = None,
        transport: httpx.BaseTransport | None = None,
        timeout: float = 30.0,
    ) -> None:
        self._api_key = api_key
        self._bearer_token = bearer_token
        self._owns_client = client is None
        self._client = client or httpx.Client(base_url=base_url, transport=transport, timeout=timeout)

    def close(self) -> None:
        if self._owns_client:
            self._client.close()

    def set_api_key(self, api_key: str | None) -> None:
        self._api_key = api_key

    def set_bearer_token(self, bearer_token: str | None) -> None:
        self._bearer_token = bearer_token

    def __enter__(self) -> GrantFlowToolClient:
        return self

    def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
        self.close()

    def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> Any:
        arguments = arguments or {}
        normalized = _unprefixed_name(name)
        if normalized == "onboard_agent":
            return self._request("POST", "/agents/onboarding", json=arguments)
        if normalized == "create_session":
            return self._request("POST", "/agents/session", json=arguments)
        if normalized == "introspect_agent":
            return self._request("POST", "/agents/introspect", json=arguments)
        if normalized == "exchange_oauth_token":
            return self._request("POST", "/agents/oauth/token", json=arguments)
        if normalized == "rotate_credential":
            return self._request("POST", "/agents/credentials/rotate", json=arguments)
        if normalized == "revoke_credential":
            return self._request("POST", "/agents/credentials/revoke", json=arguments)
        if normalized == "register_agent":
            return self._request("POST", "/agents/register", json=arguments)
        if normalized == "run_preflight":
            return self._request("POST", "/generate/preflight", json=arguments)
        if normalized == "start_generation":
            return self._request("POST", "/generate", json=arguments)
        if normalized == "get_status":
            return self._request("GET", self._job_path("/status/{job_id}", arguments))
        if normalized == "get_quality":
            return self._request("GET", self._job_path("/status/{job_id}/quality", arguments))
        if normalized == "get_events":
            return self._request("GET", self._job_path("/status/{job_id}/events", arguments))
        if normalized == "get_export_payload":
            return self._request("GET", self._job_path("/status/{job_id}/export-payload", arguments))
        if normalized == "hitl_approve":
            return self._request("POST", "/hitl/approve", json=arguments)
        if normalized == "hitl_list_pending":
            params = {k: v for k, v in arguments.items() if v is not None and k in {"tenant_id", "donor_id", "limit"}}
            return self._request("GET", "/hitl/pending", params=params or None)
        if normalized == "ingest_text":
            return self._request("POST", "/ingest/text", json=arguments)
        if normalized == "run_sandbox_happy_path":
            return self.run_sandbox_happy_path(arguments)
        raise GrantFlowToolError(
            "Unknown GrantFlow tool",
            detail={
                "code": "unknown_tool",
                "message": f"Unknown GrantFlow tool: {name}",
                "retryable": False,
                "available_tools": [tool["name"] for tool in list_tools()],
            },
        )

    def run_sandbox_happy_path(self, arguments: dict[str, Any]) -> dict[str, Any]:
        agent_name = arguments.get("agent_name")
        if not agent_name:
            raise GrantFlowToolError(
                "Missing agent_name",
                detail={
                    "code": "missing_required_argument",
                    "message": "Tool call requires agent_name.",
                    "retryable": False,
                    "required": ["agent_name"],
                },
            )
        agent_name = str(agent_name)
        registration = self.call_tool(
            "grantflow_register_agent",
            {
                "agent_name": agent_name,
                "agent_type": arguments.get("agent_type", "workflow_agent"),
                "purpose": arguments.get("purpose", "Run GrantFlow sandbox smoke workflow"),
            },
        )
        sample_requests = registration.get("sample_requests", {})
        preflight_payload = dict(sample_requests.get("preflight", {}))
        generate_payload = dict(sample_requests.get("generate", {}))
        generate_payload["request_id"] = f"{agent_name.lower().replace(' ', '-')}-{uuid.uuid4().hex[:12]}"

        preflight = self.call_tool("grantflow_run_preflight", preflight_payload)
        generation = self.call_tool("grantflow_start_generation", generate_payload)
        job_id = generation.get("job_id")
        if not job_id:
            raise GrantFlowToolError(
                "Generation response did not include job_id",
                detail={
                    "code": "missing_generation_job_id",
                    "message": "GrantFlow generation response did not include job_id.",
                    "retryable": True,
                    "generation": generation,
                },
            )
        job_id = str(job_id)
        status = self._poll_status(
            job_id,
            poll_interval_seconds=float(arguments.get("poll_interval_seconds", 0.5)),
            timeout_seconds=float(arguments.get("timeout_seconds", 20.0)),
        )
        quality = self.call_tool("grantflow_get_quality", {"job_id": job_id})
        events = self.call_tool("grantflow_get_events", {"job_id": job_id})
        return {
            "registration": registration,
            "preflight": preflight,
            "generation": generation,
            "status": status,
            "quality": quality,
            "events": events,
        }

    def _poll_status(self, job_id: str, *, poll_interval_seconds: float, timeout_seconds: float) -> Any:
        deadline = time.monotonic() + timeout_seconds
        last_status: Any = None
        while time.monotonic() <= deadline:
            last_status = self.call_tool("grantflow_get_status", {"job_id": job_id})
            if last_status.get("status") in TERMINAL_STATUSES:
                return last_status
            time.sleep(max(poll_interval_seconds, 0.0))
        raise GrantFlowToolError(
            "Timed out waiting for GrantFlow job",
            detail={
                "code": "job_poll_timeout",
                "message": f"Timed out waiting for job {job_id}",
                "retryable": True,
                "last_status": last_status,
            },
        )

    def _request(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
    ) -> Any:
        headers: dict[str, str] = {}
        if self._bearer_token:
            headers["Authorization"] = f"Bearer {self._bearer_token}"
        elif self._api_key:
            headers["X-API-Key"] = self._api_key
        if method.upper() in {"POST", "PUT", "PATCH"} and json and isinstance(json.get("request_id"), str):
            headers["X-Request-ID"] = json["request_id"]
        response = self._client.request(method, path, json=json, params=params, headers=headers or None)
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            raise GrantFlowToolError(
                "GrantFlow API request failed",
                detail={
                    "code": "grantflow_http_error",
                    "message": f"{method} {path} failed with HTTP {response.status_code}",
                    "retryable": response.status_code >= 500,
                    "status_code": response.status_code,
                    "response": _safe_json(response),
                },
            ) from exc
        return _safe_json(response)

    @staticmethod
    def _job_path(template: str, arguments: dict[str, Any]) -> str:
        job_id = arguments.get("job_id")
        if not job_id:
            raise GrantFlowToolError(
                "Missing job_id",
                detail={
                    "code": "missing_required_argument",
                    "message": "Tool call requires job_id.",
                    "retryable": False,
                    "required": ["job_id"],
                },
            )
        return template.replace("{job_id}", quote(str(job_id), safe=""))


def _safe_json(response: httpx.Response) -> Any:
    if not response.content:
        return None
    try:
        return response.json()
    except ValueError:
        return {"text": response.text}


def call_tool(
    name: str,
    arguments: dict[str, Any] | None = None,
    *,
    base_url: str = DEFAULT_BASE_URL,
    api_key: str | None = None,
    bearer_token: str | None = None,
) -> Any:
    with GrantFlowToolClient(base_url=base_url, api_key=api_key, bearer_token=bearer_token) as client:
        return client.call_tool(name, arguments)
