"""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.5.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


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._signer: Optional[Keypair] = 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._pubkey = session.pubkey  # store for claim helpers
        return session

    def authenticate(
        self,
        keypair: Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None = None,
    ) -> AgentSession:
        signer = load_keypair(keypair)
        self._signer = signer
        challenge = self.challenge()
        pubkey = str(signer.pubkey())
        self._pubkey = 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 claimable CCM: proof availability, relay readiness, recent claims."""
        pubkey = self._session_pubkey()
        return self._authorized_json(f"/v1/claims/{pubkey}/status", token=token)

    def claim(self, *, token: Optional[str] = None) -> dict[str, Any]:
        """Claim CCM via gasless relay. No SOL required — the relay pays gas."""
        pubkey = self._session_pubkey()
        return self._authorized_json(f"/v1/claims/{pubkey}/relay", method="POST", token=token)

    def deposit(
        self,
        market_id: int,
        amount_usdc: int,
        keypair: "Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None" = None,
        *,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        """Deposit USDC into a market vault via gasless relay.

        The server builds the transaction, the agent signs it locally,
        and the relay cosigns and submits. No SOL required.

        Args:
            market_id: Market ID to deposit into.
            amount_usdc: USDC amount in native units (6 decimals, e.g. 1_000_000 = 1 USDC).
            keypair: Solana keypair for signing. Uses the keypair from authenticate() if None.
        """
        import base64

        from solders.keypair import Keypair as SoldersKeypair
        from solders.transaction import Transaction

        signer = load_keypair(keypair) if keypair is not None else self._signer

        # 1. Server builds the unsigned tx (handles PDAs, ATAs, blockhash)
        build_resp = self._authorized_json(
            "/v1/agent/build-deposit",
            method="POST",
            payload={"market_id": market_id, "amount_usdc": amount_usdc},
            token=token,
        )

        # 2. Deserialize, partial-sign as user, re-serialize
        tx_bytes = base64.b64decode(build_resp["transaction"])
        tx = Transaction.from_bytes(tx_bytes)
        tx.partial_sign([signer], tx.message.recent_blockhash)
        signed_b64 = base64.b64encode(bytes(tx)).decode("ascii")

        # 3. Submit to relay for cosigning + on-chain submission
        return self._authorized_json(
            "/v1/relay/deposit",
            method="POST",
            payload={"transaction": signed_b64},
            token=token,
        )

    def contributions(self, *, limit: int = 50, offset: int = 0, token: Optional[str] = None) -> dict[str, Any]:
        """Paginated contribution history with CCM credit status."""
        return self._authorized_json(
            f"/v1/agent/contributions?limit={limit}&offset={offset}", token=token
        )

    def _session_pubkey(self) -> str:
        """Get the pubkey from the current session, or raise."""
        if not self.token:
            raise WZRDError("Not authenticated. Call authenticate() first.")
        if hasattr(self, "_pubkey") and self._pubkey:
            return self._pubkey
        raise WZRDError("Session pubkey not available. Call authenticate() first.")

    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,
        source: str = "wzrd-client",
        policy_version: str = "1.0",
        candidate_rank: Optional[int] = None,
        exploration: Optional[bool] = None,
        fallback_used: Optional[bool] = None,
        run_id: Optional[str] = None,
        cycle_id: Optional[int] = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        report_metadata = dict(metadata or {})
        wzrd_block: dict[str, Any] = {
            "source": source,
            "model": choice.model,
            "signal_model": choice.signal_model,
            "selected_model": choice.model,
            "score": choice.score,
            "raw_score": choice.raw_score,
            "trend": choice.trend,
            "confidence": choice.confidence,
            "action": choice.action,
            "platform": choice.platform,
            "policy": choice.policy,
            "policy_version": policy_version,
            "reason": choice.reason,
        }
        if candidate_rank is not None:
            wzrd_block["candidate_rank"] = candidate_rank
        if exploration is not None:
            wzrd_block["exploration"] = exploration
        if fallback_used is not None:
            wzrd_block["fallback_used"] = fallback_used
        if run_id is not None:
            wzrd_block["run_id"] = run_id
        if cycle_id is not None:
            wzrd_block["cycle_id"] = cycle_id
        report_metadata.setdefault("wzrd", wzrd_block)
        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 _env_float(name: str, default: float) -> float:
    try:
        return float(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default
