"""Async generation runner: model + dataset -> predictions -> scored report.

The async boundary is exactly "waiting on the model". Data loading and
scoring stay synchronous (fast, CPU/disk bound); this coroutine fans out
model calls with a semaphore, streams progress, then hands the collected
raw strings to the sync evaluate().

    from evalharness.model import run_eval
    report = asyncio.run(run_eval(ds, 'mock', limit=50))   # offline smoke
    report = asyncio.run(run_eval(ds, 'openai/http://gpu03:8000/v1?qwen3-8b'))
"""

import asyncio
import os
import time
from typing import Any, Dict, List, Optional, Union

from ..data.dataset import Dataset
from ..data.sample import ChatMessage, Sample
from ..eval.recipe import EvalRecipe
from ..eval.record import EvalReport
from ..eval.runner import evaluate
from .adapter import ModelAdapter, resolve_adapter
from .output import Usage


async def generate_predictions(
    adapter: ModelAdapter,
    samples: List[Sample],
    concurrency: int = 32,
    limit: Optional[int] = None,
    gen_kwargs: Optional[Dict[str, Any]] = None,
    progress: bool = True,
    env_factory=None,
    system: str = '',
    max_turns: int = 8,
    max_input_chars: int = 0,
    max_input_tokens: int = 0,
    tokenizer_path: str = '',
    attach_context_keys: tuple = ('passage', 'context'),
    limit_per_task: Optional[int] = None,
    checkpoint: Union[bool, str] = False,
    dataset_name: str = 'adhoc',
    few_shot_num: int = 0,
    few_shot_samples: Optional[List[Sample]] = None,
    few_shot_text: Optional[str] = None,
    prompt_style: str = 'strict_letter',
) -> tuple:
    """Fan out model calls; returns (pred-dicts, total_usage).

    MCQ samples are generated with the strict-letter contract ('ANSWER: X',
    evalscope parity). few_shot: official exemplar text (few_shot_text) or
    dev/train-split samples (few_shot_samples) are prepended.
    """
    gen_kwargs = gen_kwargs or {}

    def _default_max_tokens() -> int:
        return 4096

    sem = asyncio.Semaphore(concurrency)
    total_usage = Usage()
    done_count = 0
    t0 = time.time()

    def assemble(sample: Sample) -> str:
        parts = []
        if few_shot_text:
            parts.append(few_shot_text.strip())  # official exemplars, verbatim
        elif few_shot_num and few_shot_samples:
            letters_fs = 'ABCDEFGHIJ'
            for fs in few_shot_samples[:few_shot_num]:
                line = f'Question: {fs.input_text}'
                if fs.choices:
                    line += '\n' + '\n'.join(f'{letters_fs[j]}. {c}' for j, c in enumerate(fs.choices))
                ans = fs.target if not isinstance(fs.target, list) else fs.target[0]
                line += f'\nAnswer: {ans}'
                parts.append(line)
        for key in attach_context_keys:
            ctx = (sample.metadata or {}).get(key)
            if ctx:
                parts.append(str(ctx))
        question = sample.input_text
        if sample.choices:
            if prompt_style in ('strict_letter', 'auto'):
                # evalscope/OpenAI-style contract: reply ONLY 'ANSWER: X'
                letters = 'ABCDEFGHIJ'
                opts = '\n'.join(f'{letters[i]}. {c}' for i, c in enumerate(sample.choices)
                                 if i < len(letters))
                question = (f'Answer the following multiple choice question. The entire '
                            f'content of your response should be of the following format: '
                            f"'ANSWER: [LETTER]' (without quotes) where [LETTER] is one of "
                            f'{letters[:len(sample.choices)]}.\n\n{question}\n\n{opts}')
            else:
                letters = 'ABCDEFGHIJ'
                opts = '\n'.join(f'{letters[i]}. {c}' for i, c in enumerate(sample.choices)
                                 if i < len(letters))
                question = (f'{question}\n\n{opts}\n\n'
                            'Answer with the letter of the correct option.')
        elif sample.task_type in ('qa',):
            # hle OFFICIAL protocol: answer_type-specific system contract
            at = (sample.metadata or {}).get('answer_type')
            if at == 'exactMatch':
                question = (
                    'Your response should be in the following format:\n'
                    'Explanation: {your explanation for your final answer}\n'
                    'Exact Answer: {your succinct, final answer}\n'
                    'Confidence: {your confidence score between 0% and 100% for your answer}\n\n'
                    f'{question}')
            elif at == 'multipleChoice':
                question = (
                    'Your response should be in the following format:\n'
                    'Explanation: {your explanation for your answer choice}\n'
                    'Answer: {your chosen answer}\n'
                    'Confidence: {your confidence score between 0% and 100% for your answer}\n\n'
                    f'{question}')
            else:
                question = (f'{question}\n\n'
                            'End your reply with the final answer on its own last line '
                            'in the form "Answer: <answer>".')
        parts.append(question)
        text = '\n\n'.join(parts)
        if max_input_tokens:
            # reserve room for the OUTPUT budget + safety margin, else the
            # server rejects input+max_tokens > context_limit by 1 token
            budget = max(1024, max_input_tokens
                         - int(gen_kwargs.get('max_tokens') or 4096) - 2048)
            try:
                from .truncation import truncate_middle_tokens, default_tokenizer_path

                text = truncate_middle_tokens(text, budget,
                                              tokenizer_path or default_tokenizer_path())
            except Exception as e:
                # no tokenizer/transformers: degrade to a CHARS budget that
                # approximates the token cap (never send the raw 2M-token input)
                approx_chars = budget * 3
                if len(text) > approx_chars:
                    keep = approx_chars // 2
                    text = f'{text[:keep]}\n\n...[truncated {len(text) - 2 * keep} chars]...\n\n{text[-keep:]}'
                print(f'truncation degraded to chars ({type(e).__name__})', flush=True)
        if max_input_chars and len(text) > max_input_chars:
            keep = max_input_chars // 2
            head = text[:keep]
            tail = text[-keep:]
            text = f'{head}\n\n...[truncated {len(text) - 2 * keep} chars]...\n\n{tail}'
        return text

    async def one(sample: Sample) -> Dict[str, Any]:
        nonlocal done_count, total_usage
        if env_factory is not None:
            from ..agent import drive, trajectory_to_prediction
            from ..agent.loop import Environment, Usage as _U  # noqa: F401

            async with sem:
                env = env_factory()
                if type(env).run_task is not Environment.run_task:
                    # self-running env (official engine bundles: tau2/swe)
                    pred = await env.run_task(adapter, sample,
                                              max_turns=max_turns, system=system)
                    if pred is None:
                        traj = await drive(adapter, sample, env=env,
                                           max_turns=max_turns, system=system)
                        pred = trajectory_to_prediction(traj)
                else:
                    traj = await drive(adapter, sample, env=env,
                                       max_turns=max_turns, system=system)
                    pred = trajectory_to_prediction(traj)
            if not pred.get('usage'):
                pred['usage'] = traj.usage.model_dump() if 'traj' in dir() else {}
            pred.setdefault('group_key', str(sample.metadata.get('test_category')
                                             or sample.metadata.get('category')
                                             or sample.metadata.get('domain')
                                             or sample.metadata.get('id') or sample.id or ''))
            u = pred.get('usage') or {}
            total_usage = total_usage + Usage(
                input_tokens=int(u.get('input_tokens', 0) or 0),
                output_tokens=int(u.get('output_tokens', 0) or 0),
                total_tokens=int(u.get('total_tokens', 0) or 0),
                latency_s=float(u.get('latency_s', 0) or 0))
            done_count += 1
            _progress(progress, done_count, len(samples), t0, total_usage)
            return pred

        messages = ([ChatMessage(role='user', content=assemble(sample))]
                    if isinstance(sample.input, str) else list(sample.input))
        tools = None
        if sample.tools:
            tools = [{'name': t.name, 'description': t.description or '',
                      'parameters': t.parameters} for t in sample.tools]
        if getattr(adapter, 'name', '') == 'mock' \
                and adapter.extra.get('mode') in ('boxed', 'oracle', 'fc') \
                and sample.target not in ('', None):
            # oracle channel for mock verification so full pipelines run offline
            messages = messages + [ChatMessage(role='user',
                                               content=f'MOCKTARGET::{sample.target}')]
        async with sem:
            out = await adapter.generate(messages, tools=tools, **gen_kwargs)
        total_usage = total_usage + out.usage
        text = out.text
        if out.tool_calls:  # fc tasks: serialize calls as the prediction
            import json

            text = (text + '\n' if text else '') + json.dumps(
                [c.to_openai()['function'] for c in out.tool_calls], ensure_ascii=False)
        done_count += 1
        _progress(progress, done_count, len(samples), t0, total_usage)
        return {'raw': text, 'usage': out.usage.model_dump()}

    work = _apply_limits(samples, limit, limit_per_task)
    # checkpointing: restore completed samples, generate only the rest
    ckpt_store = None
    if checkpoint:
        from ..eval.checkpoint import CheckpointStore, checkpoint_path

        if isinstance(checkpoint, str):
            ckpt = checkpoint
        else:
            ckpt = checkpoint_path(os.path.expanduser('~/.cache/evalharness'),
                                   dataset_name, adapter.model or str(adapter))
        ckpt_store = CheckpointStore(ckpt, model=adapter.model or str(adapter))
        restored = ckpt_store.load()
    else:
        restored = {}

    keys = []
    pending = []
    preds_by_key: Dict[str, Dict[str, Any]] = {}
    for i, s in enumerate(work):
        k = CheckpointStore.key_for(s, i) if ckpt_store else str(i)
        keys.append(k)
        if k in restored:
            preds_by_key[k] = restored[k]
        else:
            pending.append((i, s))
    if ckpt_store is not None and restored:
        print(f'checkpoint: restored {len(restored)} predictions '
              f'({len(pending)} to generate) -> {ckpt_store.path}', flush=True)

    async def run_one(i_s):
        i, s = i_s
        pred = await one(s)
        if ckpt_store is not None:
            ckpt_store.append(keys[i], pred)
        return i, pred

    fresh = await asyncio.gather(*(run_one((i, s)) for i, s in pending))
    for i, pred in fresh:
        preds_by_key[keys[i]] = pred
    preds = [preds_by_key[k] for k in keys]
    usages = [p.get('usage', {}) for p in preds]
    return preds, usages, total_usage


def _apply_limits(samples: List[Sample], total: Optional[int],
                  per_task: Optional[int], dataset=None) -> List[Sample]:
    """total: cap the WHOLE run (ours semantics). per_task: cap each subset/
    category (evalscope's --limit semantics) -- first N per group_key."""
    if per_task:
        seen: Dict[str, int] = {}
        out = []
        for s in samples:
            key = str((s.metadata or {}).get('category')
                      or (s.metadata or {}).get('subject')
                      or (s.metadata or {}).get('test_category')
                      or getattr(getattr(dataset, 'spec', None), 'subset', 'default'))
            if seen.get(key, 0) < per_task:
                seen[key] = seen.get(key, 0) + 1
                out.append(s)
        samples = out
    if total:
        samples = samples[:total]
    return samples


def _progress(progress: bool, done: int, total: int, t0: float, usage: Usage) -> None:
    if progress and (done % 20 == 0 or done == total):
        rate = done / max(time.time() - t0, 1e-6)
        print(f'  [{done}/{total}] {rate:.1f} samples/s tokens={usage.total_tokens}', flush=True)


async def run_eval(
    dataset: Union[Dataset, List[Sample]],
    model_spec: str,
    recipe: Optional[EvalRecipe] = None,
    *,
    concurrency: int = 32,
    limit: Optional[int] = None,
    gen_kwargs: Optional[Dict[str, Any]] = None,
    judge_spec: Optional[str] = None,
    judge: Optional[Any] = None,
    progress: bool = True,
    env: str = '',
    system: str = '',
    max_turns: int = 8,
    max_input_chars: int = 0,
    max_input_tokens: int = 0,
    limit_per_task: Optional[int] = None,
    checkpoint: Union[bool, str] = False,
    dataset_name: str = 'adhoc',
    few_shot_num: int = -1,
    prompt_style: str = 'strict_letter',
) -> EvalReport:
    """Generate + score in one call. Model spec examples:
    'mock', 'mock:boxed', 'openai/http://gpu03:8000/v1?qwen3-8b', 'deploy:vllm/qwen3-8b'.

    env: environment plugin name ('bfcl_mock') -> agent message pump per
    sample; omit for single-turn generation.
    'auto' = logprob when the adapter supports it.
    few_shot_num: -1 = the dataset's declared paper default (mmlu 5, bbh 3,
    gsm8k 4, ...); 0 = zero-shot; N = explicit override.
    prompt_style: 'strict_letter' (default, evalscope-style 'ANSWER: X')
    | 'cot' (reasoning-friendly).
    """
    spec = getattr(dataset, 'spec', None)
    if few_shot_num < 0:
        few_shot_num = (spec.few_shot_num if spec is not None else 0)
    adapter = _make_adapter(model_spec)
    name = spec.name if spec is not None else 'adhoc'
    if recipe is None:
        from ..eval.recipe import EvalRecipe, get_eval

        try:
            recipe = get_eval(name)
        except KeyError:
            if name != 'adhoc':
                raise
            recipe = EvalRecipe(name='adhoc', extract='identity',
                                scorers={'acc': {'name': 'exact', 'mode': 'raw'}})
    samples = list(dataset)[:limit] if limit else list(dataset)
    samples = _apply_limits(samples, limit, limit_per_task)
    if progress:
        mode = f'agent env={env}' if env else 'single-turn'
        print(f'generating: {adapter} on {len(samples)} samples '
              f'({mode}, concurrency={concurrency})', flush=True)

    env_factory = None
    if env:
        from ..agent import ENV_REGISTRY, get_env

        if env not in ENV_REGISTRY:
            raise KeyError(f'unknown env {env!r}; available: {", ".join(ENV_REGISTRY.names())}')
        probe = get_env(env)
        if getattr(probe, 'needs_adapter', False):
            env_factory = lambda: get_env(env, adapter=adapter)  # noqa: E731
        else:
            env_factory = lambda: get_env(env)  # noqa: E731

    # paper-faithful few-shot exemplars: official hand-written hooks first
    # (bbh CoT), else the dataset's own dev/train split
    few_shot_samples = None
    few_shot_text = None
    if few_shot_num:
        from ..data.registry import get_dataset_provider

        prov = get_dataset_provider(name)
        hook = getattr(prov, 'few_shot_hook', None)
        if hook is not None:
            few_shot_text = hook('', spec.subset if spec else 'default', few_shot_num)
        if few_shot_text is None:
            fs_split = (spec.few_shot_split if spec is not None else None) or 'dev'
            try:
                import dataclasses

                fs_spec = dataclasses.replace(spec, split=fs_split) if spec is not None else None
                if fs_spec is not None:
                    from ..data.loader import load_raw_records

                    fn = prov.resolve_record_fn()
                    fs_raw = load_raw_records(fs_spec)
                    few_shot_samples = [fn(r) for r in fs_raw[:few_shot_num]]
            except Exception as e:
                print(f'few-shot: could not load {fs_split} split ({type(e).__name__}: '
                      f'{str(e)[:80]}); continuing 0-shot', flush=True)

    try:
        preds, _usages, usage = await generate_predictions(
            adapter, samples, concurrency, progress=progress,
            gen_kwargs=gen_kwargs, env_factory=env_factory,
            system=system, max_turns=max_turns, max_input_chars=max_input_chars,
            max_input_tokens=max_input_tokens,
            limit_per_task=limit_per_task,
            checkpoint=checkpoint,
            dataset_name=name,
            few_shot_num=few_shot_num,
            few_shot_samples=few_shot_samples, few_shot_text=few_shot_text,
            prompt_style=prompt_style)
    finally:
        await adapter.close()
    if judge is None and judge_spec:
        judge_adapter = _make_adapter(judge_spec)
        judge = _judge_callable(judge_adapter)

    report = evaluate(
        samples, preds, recipe,
        model=model_spec,
        judge=judge,
        extra_metadata={'gen_input_tokens': usage.input_tokens,
                        'gen_output_tokens': usage.output_tokens,
                        'gen_total_tokens': usage.total_tokens},
    )
    report.model = model_spec
    report.dataset = name
    # performance profile: pool success rate + latency/ttft percentiles
    try:
        from .aggregator import get_aggregator

        perf = get_aggregator('perf_stats')(report.samples, 'acc')
        if hasattr(adapter, 'stats'):
            perf.update({f'pool_{k}': round(v, 3) if isinstance(v, float) else v
                         for k, v in adapter.request_stats().items()})
        report.metric_groups['perf'] = perf
    except Exception:
        pass
    return report


def _make_adapter(spec: str) -> ModelAdapter:
    """Model spec forms:
    - 'mock[:mode]'                     offline adapter
    - 'openai-pool/<base-url-template>?model' with {port} placeholder:
      e.g. 'openai-pool/http://127.0.0.1:{8123..8130}/v1?Qwen3-8B' -> N ports
    - else resolve_adapter(spec)        single endpoint

    Pooled specs are CACHED per spec: all benches share one pool so the
    round-robin counter stays global (independent pools would each restart
    at the first backend and starve the rest).
    """
    from .adapter import _ADAPTER_CACHE as _CACHE

    cache_key = spec
    if cache_key in _CACHE:
        return _CACHE[cache_key]
    opts = {}
    while True:
        for f in ('!nothink', '!textools', '!perf'):
            if spec.endswith(f):
                spec = spec[:-len(f)]
                opts[f] = True
                break
        else:
            break
    if spec.startswith('openai-pool/'):
        from .pool import pooled

        rest = spec[len('openai-pool/'):]
        m = __import__('re').search(r'\{(\d+)\.\.(\d+)\}', rest)
        if not m:
            raise ValueError("openai-pool needs a {start..end} port range")
        lo, hi = int(m.group(1)), int(m.group(2))
        base_url, _, model = rest.partition('?')
        specs = []
        for port in range(lo, hi + 1):
            specs.append(f'openai/{base_url.replace(m.group(0), str(port))}?{model}')
        adapter = pooled(specs)
    elif spec.partition(':')[0] == 'mock' and ':' in spec and '/' not in spec.partition(':')[0]:
        adapter = resolve_adapter('mock')
        adapter.extra['mode'] = spec.partition(':')[2] or 'echo'
        return adapter
    else:
        adapter = resolve_adapter(spec)
    members = adapter.adapters if hasattr(adapter, 'adapters') else [adapter]
    for a in members:
        if opts.get('!nothink'):
            a.extra['no_think'] = True
        if opts.get('!textools'):
            a.extra['tools_mode'] = 'text'
        if opts.get('!perf'):
            a.extra['collect_perf'] = True
    _CACHE[cache_key] = adapter
    return adapter


def _judge_callable(judge_adapter: ModelAdapter):
    """Sync judge bridge. Works inside a running event loop (evaluate() may be
    called from async run_eval): the coroutine runs on a private loop in a
    worker thread."""

    def ask(messages) -> str:
        import asyncio

        if isinstance(messages, list) and messages and isinstance(messages[0], dict):
            messages = [ChatMessage(role=m.get('role', 'user'), content=m.get('content', ''))
                        for m in messages]

        async def go():
            out = await judge_adapter.generate(messages)
            return out.text

        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return asyncio.run(go())
        import concurrent.futures

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(asyncio.run, go()).result()

    return ask