"""Auth module — Ed25519 keypair management and token lifecycle."""

from __future__ import annotations

import base64
import json
import time
from pathlib import Path

import httpx
from nacl.signing import SigningKey


class AuthManager:
    """Manages Ed25519 keypair, registration, and token refresh.

    Credentials are cached to ``key_path`` (default:
    ``~/.fere/keys.json``).
    """

    def __init__(
        self,
        agent_name: str,
        base_url: str = "https://api.fereai.xyz",
        key_path: str | None = None,
    ) -> None:
        self.agent_name = agent_name
        self.base_url = base_url.rstrip("/")
        self._key_path = Path(key_path or Path.home() / ".fere" / "keys.json")
        self._sk: SigningKey | None = None
        self._agent_id: str | None = None
        self._token: str | None = None
        self._token_expires_at: float = 0

    # ----------------------------------------------------------
    # Keypair management
    # ----------------------------------------------------------

    def _load_or_create_keypair(self) -> None:
        """Load existing keypair from disk, or generate a new one."""
        if self._sk:
            return

        creds = self._read_creds()
        key_entry = creds.get(self.base_url, {}).get(self.agent_name)

        if key_entry and "secret_key" in key_entry:
            sk_bytes = base64.b64decode(key_entry["secret_key"])
            self._sk = SigningKey(sk_bytes)
            self._agent_id = key_entry.get("agent_id")
        else:
            self._sk = SigningKey.generate()
            self._save_keypair()

    def _save_keypair(self) -> None:
        creds = self._read_creds()
        server = creds.setdefault(self.base_url, {})
        server[self.agent_name] = {
            "secret_key": base64.b64encode(bytes(self._sk)).decode(),
            "public_key": base64.b64encode(bytes(self._sk.verify_key)).decode(),
            "agent_id": self._agent_id,
        }
        self._key_path.parent.mkdir(parents=True, exist_ok=True)
        self._key_path.write_text(json.dumps(creds, indent=2))

    def _read_creds(self) -> dict:
        if self._key_path.exists():
            return json.loads(self._key_path.read_text())
        return {}

    @property
    def public_key_b64(self) -> str:
        self._load_or_create_keypair()
        return base64.b64encode(bytes(self._sk.verify_key)).decode()

    def _sign(self, message: str) -> str:
        self._load_or_create_keypair()
        signed = self._sk.sign(message.encode())
        return base64.b64encode(signed.signature).decode()

    # ----------------------------------------------------------
    # Registration + token flow
    # ----------------------------------------------------------

    async def ensure_registered(self, client: httpx.AsyncClient) -> str:
        """Register the agent if needed. Returns agent_id."""
        self._load_or_create_keypair()
        if self._agent_id:
            return self._agent_id

        # Step 1: Register
        resp = await client.post(
            f"{self.base_url}/v1/auth/register",
            json={
                "agent_name": self.agent_name,
                "public_key": self.public_key_b64,
            },
        )
        resp.raise_for_status()
        reg = resp.json()

        # Step 2: Sign challenge and verify
        sig = self._sign(reg["challenge"])
        resp = await client.post(
            f"{self.base_url}/v1/auth/verify",
            json={
                "registration_id": reg["registration_id"],
                "challenge": reg["challenge"],
                "signature": sig,
            },
        )
        resp.raise_for_status()
        self._agent_id = resp.json()["agent_id"]
        self._save_keypair()
        return self._agent_id

    async def ensure_token(self, client: httpx.AsyncClient) -> str:
        """Get a valid bearer token, refreshing if needed."""
        if self._token and self._token_expires_at > time.time() + 30:
            return self._token

        agent_id = await self.ensure_registered(client)
        ts = str(int(time.time()))
        sig = self._sign(ts)

        resp = await client.post(
            f"{self.base_url}/v1/auth/token",
            json={
                "agent_id": agent_id,
                "timestamp": ts,
                "signature": sig,
            },
        )
        if resp.status_code == 401:
            # Agent might be re-registered — retry registration
            self._agent_id = None
            agent_id = await self.ensure_registered(client)
            ts = str(int(time.time()))
            sig = self._sign(ts)
            resp = await client.post(
                f"{self.base_url}/v1/auth/token",
                json={
                    "agent_id": agent_id,
                    "timestamp": ts,
                    "signature": sig,
                },
            )

        resp.raise_for_status()
        data = resp.json()
        self._token = data["token"]
        self._token_expires_at = time.time() + data["expires_in"]
        return self._token

    def auth_headers(self) -> dict[str, str]:
        """Return auth headers with the current token."""
        if not self._token:
            raise RuntimeError("Call ensure_token() before auth_headers()")
        return {"Authorization": f"Bearer {self._token}"}

    def clear_token(self) -> None:
        """Invalidate the cached token to force a refresh."""
        self._token = None
        self._token_expires_at = 0
