"""ModelAdapter: HOW to call a model (protocol level). Never deploys anything.

Registry + two built-ins:
  openai_compatible   -- /v1/chat/completions; covers vllm, sglang, lmdeploy,
                         ollama, tgi, and every OpenAI-protocol cloud API
  mock                -- offline deterministic/testing adapter

Model spec grammar (a single string, no config files needed for the common case):
  openai/http://localhost:8000/v1?qwen3-8b        adapter + api_base + model id
  openai/https://api.openai.com/v1?gpt-4o         (OPENAI_API_KEY read from env)
  anthropic/claude-...                            (provider-native protocols later)
  mock                                            offline
  deploy:vllm/qwen3-8b                            -> Deployer resolves to an endpoint,
                                                    then re-dispatches as openai/...
"""

import json
import os
import re
from typing import Any, Dict, List, Optional

from ..data.sample import ChatMessage
from ..eval.registry import EvalRegistry
from .output import ModelOutput, ToolCall, Usage

ADAPTER_REGISTRY = EvalRegistry('model adapter')


def register_adapter(name: str):
    def decorator(cls):
        ADAPTER_REGISTRY.register(name, cls)
        return cls

    return decorator


class ModelAdapter:
    """Base class. Subclasses implement generate() (async, structured output)."""

    name = 'base'

    def __init__(self, model: str = '', api_base: str = '', api_key: str = '', **kwargs):
        self.model = model
        self.api_base = api_base.rstrip('/')
        self.api_key = api_key
        self.extra = kwargs

    async def generate(
        self,
        messages: List[ChatMessage],
        tools: Optional[List[Dict[str, Any]]] = None,
        **gen_kwargs,
    ) -> ModelOutput:
        raise NotImplementedError

    async def score_choices(
        self,
        prompt: str,
        choices: List[str],
    ) -> 'ModelOutput':
        """Paper-faithful MCQ scoring (P1): logprob of each choice continuing
        the prompt. Strategy per backend capability:

        completions_echo  -- /completions echo+prompt_logprobs (vllm/sglang/
                            ollama); true continuation logprob, the lm-eh paper path
        chat_first_token  -- /chat/completions max_tokens=1 + top_logprobs; the
                            NEXT-token distribution encodes each choice; works on
                            every OpenAI-protocol API (cloud included); scores
                            choices as first-token probabilities

        Selected via extra['logprob_strategy'] ('auto' default).
        """
        strategy = self.extra.get('logprob_strategy', 'auto')
        if strategy in ('auto', 'completions_echo'):
            try:
                return await self._score_choices_echo(prompt, choices)
            except Exception as e:
                if strategy == 'completions_echo':
                    raise
                self.extra['logprob_strategy'] = 'chat_first_token'  # downgrade once
        return await self._score_choices_first_token(prompt, choices)

    async def _score_choices_first_token(self, prompt: str, choices: List[str]):
        """Approximate P1 via next-token top_logprobs (any chat API).

        The choices are shown in the prompt; we score P(first content token
        of each choice | prompt) from the top_logprobs of a 1-token reply.
        """
        from .output import ChoiceLogprob, ModelOutput

        letters = 'ABCDEFGHIJ'
        opts = '\n'.join(f'{letters[i]}. {c}' for i, c in enumerate(choices))
        asked = (f'{prompt}\n{opts}\n\n'
                 'Answer with only the letter of the correct option.')
        payload = {'model': self.model, 'messages': [{'role': 'user', 'content': asked}],
                   'max_tokens': 1, 'temperature': 0.0,
                   'logprobs': True, 'top_logprobs': 20}
        headers = {'Content-Type': 'application/json'}
        if self.api_key:
            headers['Authorization'] = f'Bearer {self.api_key}'
        data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
        import json as _json

        choice = (data.get('choices') or [{}])[0]
        lp_wrapper = (choice.get('logprobs') or {})
        top = lp_wrapper.get('content') or lp_wrapper.get('top_logprobs') or []
        # top: list per generated token of {token: logprob} dicts
        token_lps: Dict[str, float] = {}
        first = top[0] if top else None
        if isinstance(first, dict) and 'top_logprobs' in first:
            # OpenAI/sglang chat shape: content[0] = {token, logprob, top_logprobs:[...]}
            entries = {}
            for e in (first.get('top_logprobs') or []):
                tok, lp = e.get('token'), e.get('logprob')
                if tok is not None and lp is not None:
                    entries[str(tok)] = lp
        elif isinstance(first, dict):
            entries = first  # vllm-ish {token: lp} map
        else:
            entries = {}
        for tok, lp in entries.items():
            if isinstance(lp, dict):
                lp = lp.get('logprob')
            if lp is None:
                continue
            try:
                token_lps[str(tok).strip().upper()] = float(lp)
            except (TypeError, ValueError):
                continue
        if not token_lps:
            raise RuntimeError(f'no top_logprobs from {self.api_base}; P1 needs a '
                               'backend exposing logprobs (or use mcq_mode=generate)')
        results = []
        for i in range(len(choices)):
            letter = letters[i]
            lp = token_lps.get(letter, -100.0)
            results.append(ChoiceLogprob(text=choices[i], logprob=lp,
                                         token_logprobs=[lp], num_tokens=1))
        usage = Usage(finish_reason=choice.get('finish_reason', ''))
        return ModelOutput(model=self.model, choice_logprobs=results, usage=usage)

    async def _score_choices_echo(self, prompt: str, choices: List[str]):
        from .output import ChoiceLogprob, ModelOutput

        results: List[ChoiceLogprob] = []
        for choice in choices:
            full = f'{prompt} {choice}'
            payload = {'model': self.model, 'prompt': full, 'max_tokens': 1,
                       'echo': True, 'prompt_logprobs': 0,
                       'temperature': 0.0}
            headers = {'Content-Type': 'application/json'}
            if self.api_key:
                headers['Authorization'] = f'Bearer {self.api_key}'
            data = await self._post_raw(f'{self.api_base}/completions', payload, headers)
            # vllm/sglang shape: choices[0].prompt_logprobs: list[None|{token: lp}]
            import json as _json

            ch = (data.get('choices') or [{}])[0]
            pl = ch.get('prompt_logprobs') or []
            prompt_ids = ch.get('prompt_token_ids') or []
            # tokens belonging to the choice = the tail after the prompt part.
            # echo mode returns the whole sequence; we approximate the choice
            # token count via the tail offset encoded in logprob entries.
            lps: List[float] = []
            for entry in pl:
                if isinstance(entry, dict):
                    lp = next(iter(entry.values()))
                    if isinstance(lp, dict):
                        lp = lp.get('logprob', 0.0)
                    lps.append(float(lp))
            if not lps:
                raise RuntimeError(
                    f'backend at {self.api_base} returned no prompt_logprobs; '
                    'P1 logprob scoring needs vllm/sglang-style completions-echo'
                )
            # last len tokens we cannot know exactly; use the trailing segment
            # matching the choice's char proportion as a robust approximation
            n_prompt_chars = len(prompt) + 1
            frac = max(len(choice), 1) / max(len(full), 1)
            n_choice = max(1, round(frac * len(lps)))
            choice_lps = lps[-n_choice:]
            results.append(ChoiceLogprob(
                text=choice, logprob=sum(choice_lps),
                token_logprobs=choice_lps, num_tokens=len(choice_lps)))
        return ModelOutput(model=self.model, choice_logprobs=results)

    async def _post_raw(self, url: str, payload: Dict, headers: Dict) -> Dict[str, Any]:
        return await self._post(url, payload, headers)

    async def close(self) -> None:
        pass

    def __repr__(self):
        return f'{type(self).__name__}(model={self.model!r}, base={self.api_base!r})'


# --------------------------- spec parsing / dispatch ---------------------------

_SPEC_RE = re.compile(r'^(?P<adapter>[a-z_]+)(/(?P<rest>.+))?$')
_DEPLOY_RE = re.compile(r'^deploy:(?P<deployer>[a-z0-9_-]+)/(?P<model>.+)$')


def parse_model_spec(spec: str) -> Dict[str, str]:
    """'openai/http://x:8000/v1?qwen' -> {adapter, api_base, model}."""
    m = _SPEC_RE.match(spec.strip())
    if not m:
        raise ValueError(f'bad model spec: {spec!r} (expected adapter/rest?model)')
    adapter = m.group('adapter')
    rest = m.group('rest') or ''
    api_base, _, model = rest.partition('?')
    return {'adapter': adapter, 'api_base': api_base, 'model': model}


def resolve_adapter(spec: str, deploy_fn=None) -> ModelAdapter:
    """Spec -> adapter instance. deploy:... specs go through a Deployer first."""
    m = _DEPLOY_RE.match(spec.strip())
    if m:
        if deploy_fn is None:
            from .deployer import deploy as default_deploy
            deploy_fn = default_deploy
        endpoint = deploy_fn(m.group('deployer'), m.group('model'))
        spec = f"openai/{endpoint['api_base']}?{endpoint['model']}"
    parsed = parse_model_spec(spec)
    cls = ADAPTER_REGISTRY.get(parsed['adapter'])
    key = parsed.get('api_base') and _key_for(parsed['api_base'])
    return cls(model=parsed['model'], api_base=parsed['api_base'], api_key=key)


def _key_for(api_base: str) -> str:
    """Best-effort API key by endpoint; explicit env always wins."""
    for host_hint, var in (('api.openai.com', 'OPENAI_API_KEY'),
                           ('anthropic.com', 'ANTHROPIC_API_KEY'),
                           ('dashscope', 'DASHSCOPE_API_KEY'),
                           ('bigmodel', 'ZAI_API_KEY')):
        if host_hint in api_base:
            return os.environ.get(var, '')
    return os.environ.get('OPENAI_API_KEY', '')


# --------------------------- openai_compatible ---------------------------


@register_adapter('openai')
class OpenAICompatible(ModelAdapter):
    """Async OpenAI chat-completions client with zero hard dependencies.

    Uses httpx if installed (proper async); falls back to urllib in a thread
    so the core install stays dependency-free.
    """

    name = 'openai'

    async def generate(self, messages, tools=None, **kw) -> ModelOutput:
        payload = self._payload(messages, tools, kw)
        headers = {'Content-Type': 'application/json'}
        if self.api_key:
            headers['Authorization'] = f'Bearer {self.api_key}'
        retries = self.extra.get('retries', 3)
        last_exc: Exception = None
        for attempt in range(retries + 1):
            try:
                data = await self._post(f'{self.api_base}/chat/completions', payload, headers)
                return self._parse(data)
            except Exception as e:  # 5xx/429/timeouts: worth retrying
                last_exc = e
                retryable = 'Server error' in str(e) or '504' in str(e) or '502' in str(e) \
                    or '429' in str(e) or 'timeout' in str(e).lower() \
                    or 'TimeoutException' in type(e).__name__
                if attempt >= retries or not retryable:
                    raise
                import asyncio

                await asyncio.sleep(min(2 ** attempt * 2, 30))
        raise last_exc  # unreachable

    def _payload(self, messages, tools, kw) -> Dict[str, Any]:
        msgs = [{'role': m.role, 'content': m.content} for m in messages]
        payload: Dict[str, Any] = {'model': self.model, 'messages': msgs}
        if tools:
            payload['tools'] = [
                {'type': 'function', 'function': t} if 'function' not in t else t for t in tools
            ]
        for k in ('temperature', 'max_tokens', 'top_p', 'stop', 'seed', 'response_format'):
            if kw.get(k) is not None:
                payload[k] = kw[k]
        payload.setdefault('max_tokens', self.extra.get('max_tokens', 4096))  # CoT room
        return payload

    def _parse(self, data: Dict[str, Any]) -> ModelOutput:
        choice = (data.get('choices') or [{}])[0]
        msg = choice.get('message') or {}
        calls = []
        for c in msg.get('tool_calls') or []:
            fn = c.get('function') or {}
            args = fn.get('arguments') or '{}'
            try:
                args_dict = json.loads(args)
            except (ValueError, TypeError):
                args_dict = {}
            calls.append(ToolCall(id=c.get('id', ''), name=fn.get('name', ''),
                                  arguments=args, arguments_dict=args_dict))
        u = data.get('usage') or {}
        usage = Usage(input_tokens=u.get('prompt_tokens', 0),
                      output_tokens=u.get('completion_tokens', 0),
                      total_tokens=u.get('total_tokens', 0),
                      finish_reason=choice.get('finish_reason', ''))
        return ModelOutput(text=msg.get('content') or '', tool_calls=calls, usage=usage,
                           raw=data, model=data.get('model', self.model))

    async def _post(self, url: str, payload: Dict, headers: Dict) -> Dict[str, Any]:
        try:
            import httpx

            async with httpx.AsyncClient(timeout=self.extra.get('timeout', 600)) as client:
                r = await client.post(url, json=payload, headers=headers)
                r.raise_for_status()
                return r.json()
        except ImportError:
            import asyncio
            import urllib.request

            def _sync():
                req = urllib.request.Request(url, data=json.dumps(payload).encode(),
                                             headers=headers, method='POST')
                with urllib.request.urlopen(req, timeout=self.extra.get('timeout', 600)) as resp:
                    return json.loads(resp.read().decode())

            return await asyncio.to_thread(_sync)


# --------------------------- mock ---------------------------

_MOCK_PATTERNS = (
    (re.compile(r'\\boxed\{([^}]*)\}'), None),   # echo any boxed target in the input
)


@register_adapter('mock')
class MockAdapter(ModelAdapter):
    """Offline adapter for tests/CI/dev.

    Modes (extra['mode']):
      echo      -- return the input text (default)
      boxed     -- return \\boxed{target} (oracle channel: runner tags a
                   MOCKTARGET message when the sample carries a target)
      oracle    -- return the target verbatim (same channel; for coding
                   benches whose target is the canonical solution)
      fc        -- replay the target's ground-truth tool calls (oracle for
                   function-calling benches; target JSON in runner dict form)
      tool      -- return one tool call named extra['tool_name']
      const     -- return extra['text']
    """

    name = 'mock'

    async def generate(self, messages, tools=None, **kw) -> ModelOutput:
        mode = self.extra.get('mode', 'echo')
        if mode == 'const':
            text = self.extra.get('text', 'mock')
        elif mode in ('fc', 'tool'):
            if mode == 'tool':
                return ModelOutput(text='', tool_calls=[ToolCall(
                    name=self.extra.get('tool_name', 'dummy_tool'), arguments='{}',
                    arguments_dict={})], model='mock')
            target = None
            already_played = any(m.role == 'tool' for m in messages)
            if already_played:
                # oracle replays ground truth ONCE, then wraps up like a
                # well-behaved agent (final turn, no more calls)
                return ModelOutput(text='Done.', model='mock', usage=Usage(
                    input_tokens=1, output_tokens=1, total_tokens=2, finish_reason='stop'))
            for m in reversed(messages):
                if m.role == 'user' and m.content.startswith('MOCKTARGET::'):
                    target = m.content[len('MOCKTARGET::'):]
                    break
            calls = []
            if target:
                try:
                    gt = json.loads(target)
                    raw_calls = gt.get('tool_calls', gt if isinstance(gt, list) else [])
                    for c in raw_calls:
                        fn = c.get('function', c)
                        args = fn.get('arguments', {})
                        calls.append(ToolCall(
                            name=fn.get('name', ''), arguments=json.dumps(args),
                            arguments_dict=args if isinstance(args, dict) else {}))
                except (ValueError, TypeError):
                    calls = []
            if not calls:
                return ModelOutput(text='no tool needed', model='mock', usage=Usage(
                    input_tokens=1, output_tokens=1, total_tokens=2, finish_reason='stop'))
            return ModelOutput(text='', tool_calls=calls, model='mock', usage=Usage(
                input_tokens=1, output_tokens=len(calls), total_tokens=1 + len(calls),
                finish_reason='tool_calls'))
        else:
            last = next((m.content for m in reversed(messages) if m.role == 'user'), '')
            text = last
            target = None
            for m in reversed(messages):
                if m.role == 'user' and m.content.startswith('MOCKTARGET::'):
                    target = m.content[len('MOCKTARGET::'):]
                    break
            if mode == 'oracle':
                text = target if target is not None else last
            elif mode == 'boxed':
                if target is None:
                    m = _MOCK_PATTERNS[0][0].search(last)
                    target = m.group(1) if m else (re.findall(r'-?\d+\.?\d*', last) or ['0'])[-1]
                text = f'The answer is \\boxed{{{target}}}.'
        return ModelOutput(text=text, model='mock', usage=Usage(
            input_tokens=1, output_tokens=1, total_tokens=2, finish_reason='stop'))