"""Agent auth and reporting helpers for WZRD's CCM loop."""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping, Optional, Sequence
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from solders.keypair import Keypair

from .client import DEFAULT_TIMEOUT_SECONDS, PickResult, WZRDError

DEFAULT_API_BASE_URL = "https://api.twzrd.xyz"
DEFAULT_AGENT_USER_AGENT = "wzrd-client/0.3.0"
DEFAULT_KEYPAIR_PATH = "~/.config/solana/id.json"
_KEYPAIR_ENV_VARS = (
    "WZRD_AGENT_KEYPAIR",
    "WZRD_AGENT_KEYPAIR_PATH",
    "SOLANA_KEYPAIR",
    "SOLANA_KEYPAIR_PATH",
)


@dataclass(frozen=True)
class AgentChallenge:
    nonce: str
    message_format: str
    expires_in_secs: int


@dataclass(frozen=True)
class AgentSession:
    token: str
    pubkey: str
    expires_at: str


@dataclass(frozen=True)
class ClaimResult:
    """Result of a relay claim attempt."""

    status: str  # "success", "already_claimed", "no_proof", "rate_limited", "error"
    tx_sig: Optional[str] = None
    cumulative_total: Optional[int] = None
    retry_after_secs: Optional[int] = None
    error: Optional[str] = None


def agent_auth_message(pubkey: str, nonce: str) -> str:
    """Return the exact challenge message agents must sign."""
    return f"wzrd-agent-auth v1 | wallet:{pubkey} | nonce:{nonce}"


def load_keypair(
    keypair: Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None = None,
) -> Keypair:
    """Load a Solana keypair from common local formats.

    Supported inputs:
    - solders `Keypair`
    - path to a Solana JSON keypair file
    - JSON array string containing secret bytes
    - base58 secret string
    - raw bytes / bytearray / int sequence
    - None: checks env vars, then `~/.config/solana/id.json`
    """
    if isinstance(keypair, Keypair):
        return keypair

    if keypair is None:
        for env_name in _KEYPAIR_ENV_VARS:
            value = os.environ.get(env_name)
            if value:
                return load_keypair(value)
        default_path = Path(DEFAULT_KEYPAIR_PATH).expanduser()
        if default_path.exists():
            return load_keypair(default_path)
        raise WZRDError(
            "No Solana keypair found. Pass keypair=..., set WZRD_AGENT_KEYPAIR_PATH, "
            "or create ~/.config/solana/id.json."
        )

    if isinstance(keypair, (bytes, bytearray)):
        return Keypair.from_bytes(bytes(keypair))

    if not isinstance(keypair, (str, os.PathLike)):
        try:
            return Keypair.from_bytes(bytes(keypair))
        except (TypeError, ValueError) as exc:  # pragma: no cover - defensive
            raise WZRDError("Unsupported keypair input") from exc

    raw = str(keypair).strip()
    path = Path(raw).expanduser()
    text = path.read_text(encoding="utf-8").strip() if path.exists() else raw

    if text.startswith("["):
        try:
            secret = json.loads(text)
            return Keypair.from_bytes(bytes(secret))
        except (TypeError, ValueError, json.JSONDecodeError) as exc:
            raise WZRDError("Invalid Solana JSON keypair contents") from exc

    try:
        return Keypair.from_base58_string(text)
    except ValueError as exc:
        raise WZRDError(
            "Unsupported keypair format. Expected a Solana JSON keypair file, "
            "JSON byte array, or base58 secret string."
        ) from exc


class WZRDAgent:
    """Small synchronous helper for WZRD agent auth and contribution reporting."""

    def __init__(
        self,
        *,
        api_base_url: str = DEFAULT_API_BASE_URL,
        timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
        user_agent: str = DEFAULT_AGENT_USER_AGENT,
        token: Optional[str] = None,
    ) -> None:
        self.api_base_url = api_base_url.rstrip("/")
        self.timeout_seconds = timeout_seconds
        self.user_agent = user_agent
        self.token = token
        self._last_session: Optional[AgentSession] = None

    @classmethod
    def from_env(cls) -> "WZRDAgent":
        return cls(
            api_base_url=os.environ.get("WZRD_API_BASE_URL", DEFAULT_API_BASE_URL),
            timeout_seconds=_env_float("WZRD_TIMEOUT_SECONDS", DEFAULT_TIMEOUT_SECONDS),
            user_agent=os.environ.get("WZRD_USER_AGENT", DEFAULT_AGENT_USER_AGENT),
            token=os.environ.get("WZRD_AGENT_TOKEN"),
        )

    def challenge(self) -> AgentChallenge:
        payload = self._request_json("/v1/agent/challenge")
        return AgentChallenge(
            nonce=str(payload.get("nonce", "")),
            message_format=str(payload.get("message_format", "")),
            expires_in_secs=int(payload.get("expires_in_secs", 0) or 0),
        )

    def verify(self, *, pubkey: str, nonce: str, signature: str) -> AgentSession:
        payload = self._request_json(
            "/v1/agent/verify",
            method="POST",
            payload={
                "pubkey": pubkey,
                "nonce": nonce,
                "signature": signature,
            },
        )
        session = AgentSession(
            token=str(payload.get("token", "")),
            pubkey=str(payload.get("pubkey", pubkey)),
            expires_at=str(payload.get("expires_at", "")),
        )
        self.token = session.token
        self._last_session = session
        return session

    def authenticate(
        self,
        keypair: Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None = None,
    ) -> AgentSession:
        signer = load_keypair(keypair)
        challenge = self.challenge()
        pubkey = str(signer.pubkey())
        message = agent_auth_message(pubkey, challenge.nonce)
        signature = str(signer.sign_message(message.encode("utf-8")))
        return self.verify(pubkey=pubkey, nonce=challenge.nonce, signature=signature)

    def fund(self, *, token: Optional[str] = None) -> dict[str, Any]:
        return self._authorized_json("/v1/agent/fund", method="POST", token=token)

    def status(self, *, token: Optional[str] = None) -> dict[str, Any]:
        return self._authorized_json("/v1/agent/status", token=token)

    def earned(self, *, token: Optional[str] = None) -> dict[str, Any]:
        return self._authorized_json("/v1/agent/earned", token=token)

    def claim_status(self, *, token: Optional[str] = None) -> dict[str, Any]:
        """Check if a merkle proof exists for this agent's wallet.

        Returns proof data including cumulative_total, or raises WZRDError
        with a 404 message if no proof is available yet.
        """
        bearer = token or self.token
        if not bearer:
            raise WZRDError("Agent session token required. Call authenticate() first.")
        session = self._last_session
        if session is None:
            raise WZRDError("No session available. Call authenticate() first to set the wallet pubkey.")
        return self._authorized_json(f"/v1/claims/{session.pubkey}", token=bearer)

    def claim(self, *, token: Optional[str] = None) -> ClaimResult:
        """Claim CCM rewards via the gasless relay.

        The server pays gas — the agent never needs SOL. Flow:
        1. Check proof exists (GET /v1/claims/{pubkey})
        2. Submit relay claim (POST /v1/claims/{pubkey}/relay)

        Returns a ClaimResult with status, tx_sig, and cumulative_total.
        Per-wallet cooldown: 60s. Global limit: 120 claims/min.
        """
        bearer = token or self.token
        if not bearer:
            raise WZRDError("Agent session token required. Call authenticate() first.")
        session = self._last_session
        if session is None:
            raise WZRDError("No session available. Call authenticate() first to set the wallet pubkey.")

        pubkey = session.pubkey

        # Step 1: Check proof
        try:
            proof = self._authorized_json(f"/v1/claims/{pubkey}", token=bearer)
        except WZRDError as exc:
            msg = str(exc)
            if "404" in msg:
                return ClaimResult(status="no_proof")
            return ClaimResult(status="error", error=msg)

        if not proof or not proof.get("cumulative_total"):
            return ClaimResult(status="no_proof")

        # Step 2: Submit relay claim
        try:
            result = self._authorized_json(f"/v1/claims/{pubkey}/relay", method="POST", token=bearer)
        except WZRDError as exc:
            msg = str(exc)
            if "429" in msg:
                retry = _parse_retry_after(msg)
                return ClaimResult(status="rate_limited", retry_after_secs=retry)
            return ClaimResult(status="error", error=msg)

        # Handle already_claimed
        if result.get("status") == "already_claimed" or result.get("already_claimed"):
            return ClaimResult(
                status="already_claimed",
                cumulative_total=result.get("cumulative_total"),
            )

        return ClaimResult(
            status="success",
            tx_sig=result.get("tx_sig") or result.get("signature"),
            cumulative_total=result.get("cumulative_total"),
        )

    def report(
        self,
        *,
        model_id: str,
        task_type: Optional[str] = None,
        quality_score: Optional[float] = None,
        latency_ms: Optional[int] = None,
        cost_usd: Optional[float] = None,
        outcome: str = "success",
        metadata: Optional[Mapping[str, Any]] = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model_id": model_id,
            "outcome": outcome,
        }
        if task_type is not None:
            payload["task_type"] = task_type
        if quality_score is not None:
            payload["quality_score"] = quality_score
        if latency_ms is not None:
            payload["latency_ms"] = latency_ms
        if cost_usd is not None:
            payload["cost_usd"] = cost_usd
        if metadata is not None:
            payload["metadata"] = dict(metadata)
        return self._authorized_json("/v1/agent/report", method="POST", payload=payload, token=token)

    def report_pick(
        self,
        choice: PickResult,
        *,
        quality_score: Optional[float] = None,
        latency_ms: Optional[int] = None,
        cost_usd: Optional[float] = None,
        outcome: str = "success",
        metadata: Optional[Mapping[str, Any]] = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        report_metadata = dict(metadata or {})
        report_metadata.setdefault(
            "wzrd",
            {
                "model": choice.model,
                "signal_model": choice.signal_model,
                "score": choice.score,
                "raw_score": choice.raw_score,
                "trend": choice.trend,
                "confidence": choice.confidence,
                "action": choice.action,
                "platform": choice.platform,
                "policy": choice.policy,
                "reason": choice.reason,
            },
        )
        return self.report(
            model_id=choice.signal_model or choice.model,
            task_type=choice.task,
            quality_score=quality_score,
            latency_ms=latency_ms,
            cost_usd=cost_usd,
            outcome=outcome,
            metadata=report_metadata,
            token=token,
        )

    def _authorized_json(
        self,
        path: str,
        *,
        method: str = "GET",
        payload: Optional[Mapping[str, Any]] = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        bearer = token or self.token
        if not bearer:
            raise WZRDError("Agent session token required. Call authenticate() first or pass token=...")
        return self._request_json(path, method=method, payload=payload, token=bearer)

    def _request_json(
        self,
        path: str,
        *,
        method: str = "GET",
        payload: Optional[Mapping[str, Any]] = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        headers = {
            "Accept": "application/json",
            "User-Agent": self.user_agent,
        }
        body = None
        if token:
            headers["Authorization"] = f"Bearer {token}"
        if payload is not None:
            body = json.dumps(payload).encode("utf-8")
            headers["Content-Type"] = "application/json"

        req = Request(
            f"{self.api_base_url}{path}",
            headers=headers,
            data=body,
            method=method,
        )

        try:
            with urlopen(req, timeout=self.timeout_seconds) as resp:
                data = json.loads(resp.read())
        except HTTPError as exc:
            message = _decode_http_error(exc)
            raise WZRDError(message) from exc
        except (URLError, TimeoutError, ValueError, OSError) as exc:
            raise WZRDError(f"WZRD agent request failed: {exc}") from exc

        if not isinstance(data, dict):
            raise WZRDError("WZRD agent response was not a JSON object")
        return data


def _decode_http_error(error: HTTPError) -> str:
    try:
        payload = json.loads(error.read().decode("utf-8"))
    except (ValueError, json.JSONDecodeError, UnicodeDecodeError):
        payload = None

    if isinstance(payload, Mapping) and payload.get("error"):
        return f"WZRD agent request failed ({error.code}): {payload['error']}"
    return f"WZRD agent request failed ({error.code}): {error.reason}"


def _parse_retry_after(error_msg: str) -> Optional[int]:
    """Extract retry_after_secs from a WZRD error message containing JSON."""
    import re

    match = re.search(r"retry_after_secs[\"']?\s*[:=]\s*(\d+)", error_msg)
    if match:
        return int(match.group(1))
    return 60  # default cooldown


def _env_float(name: str, default: float) -> float:
    try:
        return float(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default
