"""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"
def _default_user_agent() -> str:
    try:
        from importlib.metadata import version as _dist_version
        return f"wzrd-client/{_dist_version('wzrd-client')}"
    except Exception:
        return "wzrd-client/1.2.3"


DEFAULT_AGENT_USER_AGENT = _default_user_agent()
SLOTS_PER_DAY = 216_000
DEFAULT_KEYPAIR_PATH = "~/.config/solana/id.json"
AUTO_KEYPAIR_PATH = "~/.config/solana/wzrd-agent.json"
_KEYPAIR_ENV_VARS = (
    "WZRD_KEYPAIR",
    "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`, then
      `~/.config/solana/wzrd-agent.json`, then auto-generates a dedicated keypair
    """
    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)
        auto_path = Path(AUTO_KEYPAIR_PATH).expanduser()
        if auto_path.exists():
            return load_keypair(auto_path)
        # Auto-generate a dedicated agent keypair
        return _auto_generate_keypair()

    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()
    if not path.exists() and ("/" in raw or "\\" in raw or raw.endswith(".json")):
        raise WZRDError(f"Keypair file not found: {path}")
    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


def _auto_generate_keypair() -> Keypair:
    """Generate a new agent keypair and save it for future sessions."""
    import logging

    logger = logging.getLogger("wzrd")
    path = Path(AUTO_KEYPAIR_PATH).expanduser()
    path.parent.mkdir(parents=True, exist_ok=True)
    kp = Keypair()
    path.write_text(json.dumps(list(bytes(kp))), encoding="utf-8")
    path.chmod(0o600)
    logger.info("generated new agent keypair: %s", kp.pubkey())
    logger.info("saved to %s", path)
    logger.info("claims are gasless — no SOL needed to earn CCM")
    return kp


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
        self._pubkey: Optional[str] = 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 invalidate_session(self) -> None:
        """Clear the cached session token, forcing re-auth on next call."""
        self.token = None

    def health_check(self) -> bool:
        """Check if the WZRD backend is healthy. Returns False on error/503."""
        try:
            resp = self._request_json("/health")
            return resp.get("status") == "ok"
        except WZRDError:
            return False

    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 infer(
        self,
        model: str,
        *,
        task_type: Optional[str] = None,
        pick_id: Optional[str] = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        """Run server-witnessed inference. Returns execution_id for verified reports.

        The server calls the LLM, grades the response, and stores an execution
        receipt. Pass the returned ``execution_id`` to :meth:`report` so the
        contribution is marked **verified** (highest reward tier).

        Args:
            model: Model identifier (must match a market channel_id).
            task_type: Task category (code, chat, reasoning, etc.).
            pick_id: Client-generated UUID for idempotency/tracing.
        """
        payload: dict[str, Any] = {"model": model}
        if task_type is not None:
            payload["task_type"] = task_type
        if pick_id is not None:
            payload["pick_id"] = pick_id
        return self._authorized_json("/v1/agent/infer", method="POST", payload=payload, token=token)

    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 ccm_balance(self) -> int:
        """Return agent's CCM balance in native units (9 decimals). Returns 0 if no ATA."""
        try:
            status = self.status()
            return status.get("ccm_balance", 0)
        except Exception:
            return 0

    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
        if signer is None:
            raise WZRDError("deposit() requires a keypair. Pass keypair= or call authenticate() first.")

        # 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 stake(
        self,
        amount: int,
        lock_duration: int = 0,
        channel_config: "str | None" = None,
        keypair: "Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None" = None,
        *,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        """Stake CCM into a channel pool via gasless relay.

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

        Args:
            amount: CCM amount in native units (9 decimals, e.g. 1_000_000_000 = 1 CCM).
            lock_duration: Lock period in slots. Use SLOTS_PER_DAY * days for convenience.
                0 = no lock (1x boost), 30 days = 1.5x, 90 days = 2x, 365 days = 3x.
            channel_config: Channel config pubkey. Uses server default if None.
            keypair: Solana keypair for signing. Uses the keypair from authenticate() if None.
        """
        import base64

        from solders.transaction import Transaction

        signer = load_keypair(keypair) if keypair is not None else self._signer
        if signer is None:
            raise WZRDError("stake() requires a keypair. Pass keypair= or call authenticate() first.")

        payload: dict[str, Any] = {"amount": amount, "lock_duration": lock_duration}
        if channel_config is not None:
            payload["channel_config"] = channel_config

        build_resp = self._authorized_json(
            "/v1/agent/build-stake",
            method="POST",
            payload=payload,
            token=token,
        )

        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")

        return self._authorized_json(
            "/v1/relay/stake",
            method="POST",
            payload={"transaction": signed_b64},
            token=token,
        )

    def unstake(
        self,
        channel_config: "str | None" = None,
        keypair: "Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None" = None,
        *,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        """Unstake CCM from a channel pool via gasless relay. Lock must be expired."""
        import base64

        from solders.transaction import Transaction

        signer = load_keypair(keypair) if keypair is not None else self._signer
        if signer is None:
            raise WZRDError("unstake() requires a keypair. Pass keypair= or call authenticate() first.")

        payload: dict[str, Any] = {}
        if channel_config is not None:
            payload["channel_config"] = channel_config

        build_resp = self._authorized_json(
            "/v1/agent/build-unstake",
            method="POST",
            payload=payload,
            token=token,
        )

        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")

        return self._authorized_json(
            "/v1/relay/unstake",
            method="POST",
            payload={"transaction": signed_b64},
            token=token,
        )

    def claim_staking_rewards(
        self,
        channel_config: "str | None" = None,
        keypair: "Keypair | str | os.PathLike[str] | Sequence[int] | bytes | bytearray | None" = None,
        *,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        """Claim pending staking rewards via gasless relay."""
        import base64

        from solders.transaction import Transaction

        signer = load_keypair(keypair) if keypair is not None else self._signer
        if signer is None:
            raise WZRDError("claim_staking_rewards() requires a keypair. Pass keypair= or call authenticate() first.")

        payload: dict[str, Any] = {}
        if channel_config is not None:
            payload["channel_config"] = channel_config

        build_resp = self._authorized_json(
            "/v1/agent/build-claim-rewards",
            method="POST",
            payload=payload,
            token=token,
        )

        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")

        return self._authorized_json(
            "/v1/relay/claim-rewards",
            method="POST",
            payload={"transaction": signed_b64},
            token=token,
        )

    def staking_status(
        self,
        *,
        channel_config: "str | None" = None,
        token: Optional[str] = None,
    ) -> dict[str, Any]:
        """Check staking position: amount, lock, rewards, pool stats."""
        path = "/v1/agent/staking"
        if channel_config is not None:
            path += f"?channel_config={channel_config}"
        return self._authorized_json(path, 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,
        execution_id: 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 execution_id is not None:
            payload["execution_id"] = execution_id
        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,
        *,
        execution_id: 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,
        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,
            execution_id=execution_id,
            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, status_code=exc.code) 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
