"""
Elysium — Connect your ADK agent to AitherOS cloud inference.
================================================================

The Elysium client is the onramp for external agents to register with
the AitherOS platform and use its inference pipeline, scoped to their
tenant's plan and token balance.

Quick start::

    from adk import AitherAgent
    from adk.elysium import Elysium

    # Connect to Elysium (authenticates + registers your agent)
    elysium = await Elysium.connect(api_key="aither_sk_live_...")

    # Create an agent that uses Elysium inference
    agent = AitherAgent("my_agent", llm=elysium.router)
    response = await agent.chat("Hello from my custom agent!")

    # Or connect an existing agent
    agent = AitherAgent("my_agent")
    await elysium.attach(agent)
    response = await agent.chat("Now using Elysium inference!")

Authentication::

    # Register a new account
    elysium = Elysium()
    result = await elysium.register(email="you@example.com", password="...")

    # Login to get a JWT
    result = await elysium.login(email="you@example.com", password="...")
    # result["token"] is your session token

    # Or use an ACTA API key (from demo.aitherium.com)
    elysium = await Elysium.connect(api_key="aither_sk_live_xxx")

All inference is routed through mcp.aitherium.com/v1/chat/completions
with ACTA token metering.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional

import httpx

if TYPE_CHECKING:
    from adk.agent import AitherAgent
    from adk.llm import LLMRouter

logger = logging.getLogger("adk.elysium")

# Elysium endpoints
_GATEWAY_URL = "https://gateway.aitherium.com"
_INFERENCE_URL = "https://mcp.aitherium.com/v1"


@dataclass
class ElysiumStatus:
    """Connection status for an Elysium-connected agent."""
    connected: bool = False
    user_id: str = ""
    plan: str = ""
    tier: str = ""
    token_balance: int = 0
    models_available: list = field(default_factory=list)
    agent_id: str = ""


class Elysium:
    """Client for connecting ADK agents to AitherOS Elysium cloud inference.

    Provides:
    - User registration and authentication (gateway.aitherium.com)
    - Agent registration with the network
    - Cloud inference via mcp.aitherium.com (OpenAI-compatible)
    - Model discovery and tier-based access
    - ACTA token balance checking

    All features are optional — agents work fully offline without Elysium.
    """

    def __init__(
        self,
        api_key: str = "",
        gateway_url: str = "",
        inference_url: str = "",
        timeout: float = 30.0,
    ):
        self.api_key = api_key or os.getenv("AITHER_API_KEY", "")
        self.gateway_url = (gateway_url or os.getenv("AITHER_GATEWAY_URL", _GATEWAY_URL)).rstrip("/")
        self.inference_url = (inference_url or os.getenv("AITHER_INFERENCE_URL", _INFERENCE_URL)).rstrip("/")
        self._timeout = timeout
        self._status = ElysiumStatus()
        self._router: Optional[LLMRouter] = None
        self._jwt_token: str = ""  # Session token from login

    def _auth_headers(self) -> dict[str, str]:
        """Build auth headers using API key or JWT token."""
        h: dict[str, str] = {"Content-Type": "application/json"}
        token = self.api_key or self._jwt_token
        if token:
            h["Authorization"] = f"Bearer {token}"
        return h

    def _gateway_client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(
            timeout=self._timeout,
            headers={"Content-Type": "application/json"},
        )

    def _inference_client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(
            timeout=self._timeout,
            headers=self._auth_headers(),
        )

    # ─── Authentication ───────────────────────────────────────────────

    async def register(self, email: str, password: str, display_name: str = "") -> dict:
        """Register a new Aitherium account.

        Returns {"ok": True, "user_id": "..."} on success.
        After registration, verify your email, then login to get a JWT.
        """
        async with self._gateway_client() as client:
            resp = await client.post(
                f"{self.gateway_url}/v1/auth/register",
                json={
                    "email": email,
                    "password": password,
                    "display_name": display_name or None,
                },
            )
            resp.raise_for_status()
            return resp.json()

    async def login(self, email: str, password: str) -> dict:
        """Login to get a JWT session token.

        Returns {"ok": True, "token": "...", "expires_at": ...}.
        The token is automatically stored for subsequent API calls.
        """
        async with self._gateway_client() as client:
            resp = await client.post(
                f"{self.gateway_url}/v1/auth/login",
                json={"email": email, "password": password},
            )
            resp.raise_for_status()
            data = resp.json()

        if data.get("ok") and data.get("token"):
            self._jwt_token = data["token"]
            logger.info("Logged in to Aitherium as %s", email)

        return data

    async def verify_email(self, email: str, token: str) -> dict:
        """Verify email with the token from the verification email."""
        async with self._gateway_client() as client:
            resp = await client.post(
                f"{self.gateway_url}/v1/auth/verify",
                json={"email": email, "token": token},
            )
            resp.raise_for_status()
            return resp.json()

    # ─── Connection ───────────────────────────────────────────────────

    @classmethod
    async def connect(
        cls,
        api_key: str = "",
        gateway_url: str = "",
        inference_url: str = "",
    ) -> "Elysium":
        """Connect to Elysium and verify credentials.

        Args:
            api_key: ACTA API key (aither_sk_live_...) or JWT token.
            gateway_url: Override gateway URL (default: gateway.aitherium.com).
            inference_url: Override inference URL (default: mcp.aitherium.com/v1).

        Returns:
            Connected Elysium instance with a configured LLM router.

        Raises:
            ConnectionError: If credentials are invalid or gateway unreachable.
        """
        instance = cls(
            api_key=api_key,
            gateway_url=gateway_url,
            inference_url=inference_url,
        )
        await instance._verify_connection()
        return instance

    async def _verify_connection(self) -> None:
        """Verify credentials and fetch available models."""
        if not self.api_key and not self._jwt_token:
            raise ConnectionError(
                "No API key or JWT token provided.\n\n"
                "  Get an API key: https://demo.aitherium.com\n"
                "  Or login:       await elysium.login(email, password)\n"
                "  Or set env:     AITHER_API_KEY=aither_sk_live_..."
            )

        # Check health
        try:
            async with self._inference_client() as client:
                resp = await client.get(f"{self.inference_url.rsplit('/v1', 1)[0]}/health")
                if resp.status_code != 200:
                    raise ConnectionError(f"Elysium health check failed: {resp.status_code}")
        except httpx.ConnectError:
            raise ConnectionError(
                "Cannot reach mcp.aitherium.com. Check your internet connection."
            )

        # Fetch available models (validates auth implicitly via the inference endpoint)
        try:
            async with self._inference_client() as client:
                resp = await client.get(f"{self.inference_url}/models")
                if resp.status_code == 200:
                    data = resp.json()
                    self._status.models_available = [
                        m["id"] for m in data.get("data", [])
                        if m.get("accessible", True)
                    ]
        except Exception:
            pass  # Models endpoint is public, non-fatal

        self._status.connected = True
        logger.info(
            "Connected to Elysium (%s). Models: %s",
            self.inference_url,
            ", ".join(self._status.models_available) or "discovering...",
        )

    @property
    def router(self) -> "LLMRouter":
        """Get an LLMRouter configured for Elysium inference.

        This router points at mcp.aitherium.com/v1 with your API key,
        speaking standard OpenAI-compatible format.
        """
        if self._router is None:
            from adk.llm import LLMRouter
            self._router = LLMRouter(
                provider="gateway",
                base_url=self.inference_url,
                api_key=self.api_key or self._jwt_token,
                model="aither-orchestrator",
            )
        return self._router

    @property
    def status(self) -> ElysiumStatus:
        """Current connection status."""
        return self._status

    # ─── Agent Integration ────────────────────────────────────────────

    async def attach(self, agent: "AitherAgent") -> None:
        """Attach Elysium inference to an existing agent.

        Replaces the agent's LLM router with one pointing at Elysium.
        The agent's custom tools, identity, and memory are preserved —
        only inference is redirected to the cloud.
        """
        if not self._status.connected:
            await self._verify_connection()

        agent.llm = self.router
        agent._provider_name = "elysium"
        logger.info("Attached Elysium inference to agent '%s'", agent.name)

    # ─── Agent Registry ───────────────────────────────────────────────

    async def register_agent(
        self,
        agent_name: str,
        capabilities: list[str] | None = None,
        description: str = "",
        tools: list[str] | None = None,
    ) -> dict:
        """Register your agent with the Aitherium network.

        Makes your agent discoverable by others on the platform.
        """
        async with httpx.AsyncClient(
            timeout=self._timeout,
            headers=self._auth_headers(),
        ) as client:
            resp = await client.post(
                f"{self.gateway_url}/v1/agents/register",
                json={
                    "name": agent_name,
                    "capabilities": capabilities or [],
                    "description": description,
                    "tools": tools or [],
                },
            )
            resp.raise_for_status()
            data = resp.json()

        self._status.agent_id = data.get("agent_id", "")
        return data

    async def my_agents(self) -> list[dict]:
        """List your registered agents."""
        async with httpx.AsyncClient(
            timeout=self._timeout,
            headers=self._auth_headers(),
        ) as client:
            resp = await client.get(f"{self.gateway_url}/v1/agents/mine")
            if resp.status_code == 404:
                # Endpoint may not exist yet on older gateways
                return []
            resp.raise_for_status()
            return resp.json().get("agents", [])

    async def discover_agents(
        self,
        capability: str | None = None,
        limit: int = 20,
    ) -> list[dict]:
        """Discover other agents on the Aitherium network."""
        params: dict = {"limit": limit}
        if capability:
            params["capability"] = capability
        async with httpx.AsyncClient(
            timeout=self._timeout,
            headers=self._auth_headers(),
        ) as client:
            resp = await client.get(
                f"{self.gateway_url}/v1/agents/discover",
                params=params,
            )
            resp.raise_for_status()
            return resp.json().get("agents", [])

    # ─── Utilities ────────────────────────────────────────────────────

    async def models(self) -> list[dict]:
        """List available inference models with tier and pricing info."""
        async with self._inference_client() as client:
            resp = await client.get(f"{self.inference_url}/models")
            resp.raise_for_status()
            return resp.json().get("data", [])

    async def health(self) -> bool:
        """Check if Elysium is reachable."""
        try:
            async with self._inference_client() as client:
                base = self.inference_url.rsplit("/v1", 1)[0]
                resp = await client.get(f"{base}/health")
                return resp.status_code == 200
        except Exception:
            return False

    async def preflight(self) -> dict:
        """Run comprehensive health checks for agent self-diagnosis.

        Returns a dict with check results that agents can use to validate
        their setup before starting operations::

            status = await ely.preflight()
            if status["ready"]:
                # All systems go
            else:
                # Check status["checks"] for failures

        Checks performed:
        - gateway_reachable: Can we reach gateway.aitherium.com?
        - inference_reachable: Can we reach mcp.aitherium.com?
        - auth_valid: Is our API key / JWT accepted?
        - balance: Current token balance and plan info
        - models: Available models for our tier
        - agents: Our registered agents
        """
        checks: dict = {
            "ready": False,
            "checks": {},
        }

        # 1. Gateway health
        try:
            async with self._gateway_client() as client:
                resp = await client.get(f"{self.gateway_url}/health")
                checks["checks"]["gateway_reachable"] = {
                    "ok": resp.status_code == 200,
                    "status_code": resp.status_code,
                    "url": f"{self.gateway_url}/health",
                }
        except Exception as e:
            checks["checks"]["gateway_reachable"] = {
                "ok": False, "error": str(e),
                "url": f"{self.gateway_url}/health",
            }

        # 2. Inference health
        try:
            async with self._inference_client() as client:
                base = self.inference_url.rsplit("/v1", 1)[0]
                resp = await client.get(f"{base}/health")
                checks["checks"]["inference_reachable"] = {
                    "ok": resp.status_code == 200,
                    "status_code": resp.status_code,
                    "url": f"{base}/health",
                }
        except Exception as e:
            checks["checks"]["inference_reachable"] = {
                "ok": False, "error": str(e),
                "url": self.inference_url,
            }

        # 3. Auth validation + balance (via billing endpoint)
        try:
            async with httpx.AsyncClient(
                timeout=self._timeout, headers=self._auth_headers()
            ) as client:
                resp = await client.get(f"{self.gateway_url}/v1/billing/balance")
                if resp.status_code == 200:
                    data = resp.json()
                    checks["checks"]["auth_valid"] = {"ok": True}
                    checks["checks"]["balance"] = {
                        "ok": True,
                        "balance": data.get("balance", 0),
                        "plan": data.get("plan", "unknown"),
                    }
                elif resp.status_code == 401:
                    checks["checks"]["auth_valid"] = {
                        "ok": False, "error": "Invalid API key or token",
                    }
                    checks["checks"]["balance"] = {"ok": False, "error": "Auth failed"}
                else:
                    checks["checks"]["auth_valid"] = {
                        "ok": False, "status_code": resp.status_code,
                    }
                    checks["checks"]["balance"] = {"ok": False, "status_code": resp.status_code}
        except Exception as e:
            checks["checks"]["auth_valid"] = {"ok": False, "error": str(e)}
            checks["checks"]["balance"] = {"ok": False, "error": str(e)}

        # 4. Models
        try:
            async with self._inference_client() as client:
                resp = await client.get(f"{self.inference_url}/models")
                if resp.status_code == 200:
                    data = resp.json()
                    model_list = data.get("data", [])
                    accessible = [m["id"] for m in model_list if m.get("accessible", True)]
                    checks["checks"]["models"] = {
                        "ok": len(accessible) > 0,
                        "total": len(model_list),
                        "accessible": accessible,
                    }
                else:
                    checks["checks"]["models"] = {"ok": False, "status_code": resp.status_code}
        except Exception as e:
            checks["checks"]["models"] = {"ok": False, "error": str(e)}

        # 5. Registered agents
        try:
            agents = await self.my_agents()
            checks["checks"]["agents"] = {
                "ok": True,
                "count": len(agents),
                "names": [a.get("agent_name", a.get("name", "")) for a in agents],
            }
        except Exception as e:
            checks["checks"]["agents"] = {"ok": False, "error": str(e)}

        # Overall readiness
        critical = ["gateway_reachable", "inference_reachable"]
        checks["ready"] = all(
            checks["checks"].get(c, {}).get("ok", False) for c in critical
        )

        return checks
