"""WZRD Router — wraps chat clients with momentum-informed model selection.

Usage:
    import openai
    from wzrd.router import WZRDRouter

    client = WZRDRouter(openai.Client())
    response = client.chat.completions.create(
        model="code",  # task sentinel — WZRD picks the actual model
        messages=[{"role": "user", "content": "Write a Python function"}],
    )
    # response.model tells you which model WZRD chose

Auto-reporting (opt-in):
    from wzrd.router import WZRDRouter
    from wzrd.agent import WZRDAgent

    agent = WZRDAgent()
    agent.authenticate("~/.config/solana/id.json")

    client = WZRDRouter(openai.Client(), report=True, agent=agent)
    # Every routed call now auto-reports for WZRD rewards.

    # Or from env vars (WZRD_REPORT=true, WZRD_AGENT_KEYPAIR_PATH=...):
    client = WZRDRouter.from_env(openai.Client())

Works with clients that expose client.chat.completions.create().
Explicit model names pass through unchanged by default; only task sentinels
or missing model values trigger routing.
"""

from __future__ import annotations

import logging
import os
import threading
from typing import Any, Optional

from .client import PickResult, pick_details

logger = logging.getLogger("wzrd.router")


def _fire_and_forget_report(agent: Any, choice: PickResult) -> None:
    """Report a routing decision in a background thread. Never raises."""
    try:
        agent.report_pick(choice)
    except Exception:
        logger.warning("wzrd.router: auto-report failed for %s", choice.model, exc_info=True)


class _ChatCompletions:
    """Proxy for client.chat.completions that injects WZRD model selection."""

    def __init__(
        self,
        real_completions: Any,
        task: str,
        fallback: str | None,
        candidates: list[str] | None,
        agent: Any,
        report: bool,
    ):
        self._real = real_completions
        self._task = task
        self._fallback = fallback
        self._candidates = candidates
        self._agent = agent
        self._report = report

    # Task sentinels that trigger WZRD routing. Anything else passes through.
    _TASK_SENTINELS = {
        "code", "chat", "reasoning", "general", "math", "coding", "conversation",
        "vision", "image", "audio", "tts", "embedding",
    }

    def _resolve_model(self, model: str | None) -> tuple[str, Optional[PickResult]]:
        """Resolve the model parameter using WZRD's prior.

        Returns (model_name, pick_result_or_None).

        - model=None -> use WZRD with default task
        - model="code"/"chat"/etc -> WZRD picks for that task
        - model="anthropic/claude-sonnet-4.6" -> pass through unchanged
        """
        if model is None:
            # No model specified -- WZRD picks
            task = self._task
        elif model in self._TASK_SENTINELS:
            # Task sentinel -- WZRD picks for this task
            task = model
        else:
            # Explicit model name -- pass through, don't override
            return model, None

        result = pick_details(task, candidates=self._candidates, fallback=self._fallback)
        if result and result.model:
            logger.info("wzrd.router: %s -> %s", task, result.model)
            return result.model, result
        fallback_model = self._fallback or model or "gpt-4o"
        return fallback_model, None

    def create(self, **kwargs: Any) -> Any:
        chosen, pick_result = self._resolve_model(kwargs.get("model"))
        kwargs["model"] = chosen
        response = self._real.create(**kwargs)
        if self._report and pick_result is not None and self._agent is not None:
            thread = threading.Thread(
                target=_fire_and_forget_report,
                args=(self._agent, pick_result),
                daemon=True,
            )
            thread.start()
        return response

    async def acreate(self, **kwargs: Any) -> Any:
        chosen, pick_result = self._resolve_model(kwargs.get("model"))
        kwargs["model"] = chosen
        if hasattr(self._real, "acreate"):
            response = await self._real.acreate(**kwargs)
        else:
            response = self._real.create(**kwargs)
        if self._report and pick_result is not None and self._agent is not None:
            thread = threading.Thread(
                target=_fire_and_forget_report,
                args=(self._agent, pick_result),
                daemon=True,
            )
            thread.start()
        return response

    def __getattr__(self, name: str) -> Any:
        return getattr(self._real, name)


class _Chat:
    """Proxy for client.chat that wraps completions."""

    def __init__(
        self,
        real_chat: Any,
        task: str,
        fallback: str | None,
        candidates: list[str] | None,
        agent: Any,
        report: bool,
    ):
        self._real = real_chat
        self._task = task
        self._fallback = fallback
        self._candidates = candidates
        self._agent = agent
        self._report = report

    @property
    def completions(self) -> _ChatCompletions:
        return _ChatCompletions(
            self._real.completions,
            self._task,
            self._fallback,
            self._candidates,
            self._agent,
            self._report,
        )

    def __getattr__(self, name: str) -> Any:
        return getattr(self._real, name)


class WZRDRouter:
    """Wrap a chat client with WZRD momentum-informed model selection.

    Args:
        client: Any client exposing client.chat.completions.create(...)
        task: Default task hint for model selection ("code", "chat", "general")
        fallback: Model to use if WZRD is unavailable
        candidates: Optional allowlist of model names to consider
        report: If True, auto-report every routing decision via WZRDAgent (default False)
        agent: Pre-authenticated WZRDAgent instance for reporting
        keypair: Solana keypair (path, Keypair, or bytes) for auto-authenticating a new agent

    When ``report=True``, one of ``agent`` or ``keypair`` must be provided.
    If ``keypair`` is given without ``agent``, a new WZRDAgent is created and
    authenticated automatically.

    Example:
        import openai
        from wzrd.router import WZRDRouter

        # WZRD picks the model when you omit model or pass a task sentinel
        client = WZRDRouter(openai.Client(), task="code")
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": "hello"}],
        )

        # Auto-report for rewards:
        from wzrd.agent import WZRDAgent
        agent = WZRDAgent()
        agent.authenticate("~/.config/solana/id.json")
        client = WZRDRouter(openai.Client(), report=True, agent=agent)

        # Or from environment:
        client = WZRDRouter.from_env(openai.Client())
    """

    def __init__(
        self,
        client: Any,
        *,
        task: str = "general",
        fallback: str | None = None,
        candidates: list[str] | None = None,
        report: bool = False,
        agent: Any = None,
        keypair: Any = None,
    ):
        self._client = client
        self._task = task
        self._fallback = fallback
        self._candidates = candidates
        self._report = report
        self._agent = agent

        if report and agent is None and keypair is not None:
            self._agent = self._make_agent(keypair)

    @classmethod
    def from_env(
        cls,
        client: Any,
        *,
        task: str = "general",
        fallback: str | None = None,
        candidates: list[str] | None = None,
    ) -> "WZRDRouter":
        """Create a WZRDRouter configured from environment variables.

        Reads:
            WZRD_REPORT: "true" / "1" to enable auto-reporting (default: off)
            WZRD_AGENT_KEYPAIR_PATH: path to Solana keypair for agent auth
            WZRD_AGENT_KEYPAIR: inline keypair (JSON array or base58)
        """
        report = os.environ.get("WZRD_REPORT", "").lower() in ("true", "1", "yes")
        keypair_val: str | None = None
        for env_name in ("WZRD_AGENT_KEYPAIR_PATH", "WZRD_AGENT_KEYPAIR"):
            val = os.environ.get(env_name)
            if val:
                keypair_val = val
                break
        return cls(
            client,
            task=task,
            fallback=fallback,
            candidates=candidates,
            report=report,
            keypair=keypair_val if report else None,
        )

    @staticmethod
    def _make_agent(keypair: Any) -> Any:
        """Create and authenticate a WZRDAgent from a keypair. Returns None on failure."""
        try:
            from .agent import WZRDAgent

            agent = WZRDAgent.from_env()
            agent.authenticate(keypair)
            return agent
        except Exception:
            logger.warning("wzrd.router: agent authentication failed; reporting disabled", exc_info=True)
            return None

    @property
    def chat(self) -> _Chat:
        real_chat = getattr(self._client, "chat", None)
        completions = getattr(real_chat, "completions", None) if real_chat is not None else None
        create = getattr(completions, "create", None) if completions is not None else None
        if real_chat is None or completions is None or not callable(create):
            raise TypeError(
                "WZRDRouter requires a client exposing client.chat.completions.create(...)"
            )
        return _Chat(real_chat, self._task, self._fallback, self._candidates, self._agent, self._report)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)
