"""OAuth 2.0 Device Authorization Grant (RFC 8628) client.

Authenticates the CLI with LLMHosts.com cloud:
1. Request device code from server
2. Display user code + verification URL
3. Poll until user approves in browser
4. Store credentials locally

Usage::

    flow = DeviceCodeFlow(server_url="https://llmhosts.com")
    creds = await flow.authenticate()
    print(f"API key: {creds.api_key}")
"""

from __future__ import annotations

import asyncio
import json
import logging
import webbrowser
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable

import httpx

logger = logging.getLogger(__name__)

# Default config directory
_CONFIG_DIR = Path.home() / ".llmhosts"
_CREDENTIALS_FILE = "credentials.json"


@dataclass
class Credentials:
    """Stored authentication credentials."""

    api_key: str
    user_id: str
    email: str
    name: str | None
    plan: str
    server_url: str

    def to_dict(self) -> dict[str, str | None]:
        return {
            "api_key": self.api_key,
            "user_id": self.user_id,
            "email": self.email,
            "name": self.name,
            "plan": self.plan,
            "server_url": self.server_url,
        }

    @classmethod
    def from_dict(cls, data: dict[str, str | None]) -> Credentials:
        return cls(
            api_key=str(data["api_key"]),
            user_id=str(data["user_id"]),
            email=str(data["email"]),
            name=data.get("name"),
            plan=str(data.get("plan", "free")),
            server_url=str(data.get("server_url", "https://llmhosts.com")),
        )


class DeviceCodeFlow:
    """RFC 8628 Device Authorization Grant client."""

    def __init__(
        self,
        server_url: str = "https://llmhosts.com",
        config_dir: Path | None = None,
    ) -> None:
        self._server_url = server_url.rstrip("/")
        self._config_dir = config_dir or _CONFIG_DIR
        self._credentials_path = self._config_dir / _CREDENTIALS_FILE

    async def authenticate(
        self,
        *,
        auto_open_browser: bool = True,
        on_user_code: Callable[[str, str], None] | None = None,
    ) -> Credentials:
        """Run the full device code flow.

        Parameters
        ----------
        auto_open_browser:
            Attempt to open the verification URL in the user's browser.
        on_user_code:
            Callback called with (user_code, verification_uri) for custom display.

        Returns
        -------
        Credentials
            The authenticated credentials with API key.

        Raises
        ------
        AuthenticationError
            If the flow fails (expired, network error, etc.).
        """
        async with httpx.AsyncClient(timeout=30.0) as client:
            # Step 1: Request device code
            resp = await client.post(f"{self._server_url}/api/auth/device-code")
            if resp.status_code != 200:
                raise AuthenticationError(f"Failed to get device code: {resp.text}")

            data = resp.json()
            device_code = data["device_code"]
            user_code = data["user_code"]
            verification_uri = data["verification_uri"]
            expires_in = data.get("expires_in", 900)
            interval = data.get("interval", 5)

            # Step 2: Display code
            if on_user_code:
                on_user_code(user_code, verification_uri)
            else:
                logger.info("Visit %s and enter code: %s", verification_uri, user_code)

            # Step 3: Open browser
            if auto_open_browser:
                import contextlib

                with contextlib.suppress(Exception):
                    webbrowser.open(verification_uri)

            # Step 4: Poll until approved or expired
            deadline = asyncio.get_event_loop().time() + expires_in
            while asyncio.get_event_loop().time() < deadline:
                await asyncio.sleep(interval)

                poll_resp = await client.post(
                    f"{self._server_url}/api/auth/device-code/poll",
                    json={"device_code": device_code},
                )

                if poll_resp.status_code == 200:
                    # Approved!
                    token_data = poll_resp.json()
                    creds = Credentials(
                        api_key=token_data["api_key"],
                        user_id=token_data["user"]["id"],
                        email=token_data["user"]["email"],
                        name=token_data["user"].get("name"),
                        plan=token_data["user"].get("plan", "free"),
                        server_url=self._server_url,
                    )
                    await self._save_credentials(creds)
                    return creds

                if poll_resp.status_code == 428:
                    # authorization_pending -- keep polling
                    continue

                if poll_resp.status_code == 410:
                    raise AuthenticationError("Device code expired. Please try again.")

                # Unexpected error
                error_data = poll_resp.json()
                raise AuthenticationError(
                    f"Authentication failed: {error_data.get('error_description', poll_resp.text)}"
                )

            raise AuthenticationError("Authentication timed out. Please try again.")

    async def load_credentials(self) -> Credentials | None:
        """Load saved credentials from disk. Returns None if not found."""
        if not self._credentials_path.is_file():
            return None
        try:
            data = json.loads(self._credentials_path.read_text(encoding="utf-8"))
            return Credentials.from_dict(data)
        except (json.JSONDecodeError, KeyError, TypeError) as exc:
            logger.warning("Failed to load credentials: %s", exc)
            return None

    async def refresh_plan(self, creds: Credentials) -> Credentials:
        """Fetch the live plan from the server and persist updated credentials.

        Returns the updated Credentials (plan field reflects current subscription).
        On any network / auth error the original credentials are returned unchanged.
        """
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(
                    f"{self._server_url}/api/auth/me",
                    headers={"Authorization": f"Bearer {creds.api_key}"},
                )
            if resp.status_code == 200:
                data = resp.json()
                updated = Credentials(
                    api_key=creds.api_key,
                    user_id=data.get("id", creds.user_id),
                    email=data.get("email", creds.email),
                    name=data.get("name", creds.name),
                    plan=data.get("plan", creds.plan),
                    server_url=creds.server_url,
                )
                if updated.plan != creds.plan:
                    logger.info("Plan updated from %s → %s", creds.plan, updated.plan)
                    await self._save_credentials(updated)
                return updated
        except Exception as exc:
            logger.debug("Plan refresh failed (non-fatal): %s", exc)
        return creds

    async def clear_credentials(self) -> None:
        """Remove stored credentials."""
        if self._credentials_path.is_file():
            self._credentials_path.unlink()
            logger.info("Credentials cleared")

    async def _save_credentials(self, creds: Credentials) -> None:
        """Save credentials to disk."""
        self._config_dir.mkdir(parents=True, exist_ok=True)
        self._credentials_path.write_text(
            json.dumps(creds.to_dict(), indent=2),
            encoding="utf-8",
        )
        # Restrict permissions (owner only)
        self._credentials_path.chmod(0o600)
        logger.info("Credentials saved to %s", self._credentials_path)


class AuthenticationError(Exception):
    """Raised when device code authentication fails."""
