"""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 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,
    attach_context_keys: tuple = ('passage', 'context'),
    limit_per_task: Optional[int] = None,
    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 {}
    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',):
            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_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

            async with sem:
                traj = await drive(adapter, sample, env=env_factory(),
                                   max_turns=max_turns, system=system)
            total_usage = total_usage + traj.usage
            pred = trajectory_to_prediction(traj)
            pred['group_key'] = str(sample.metadata.get('test_category')
                                    or sample.metadata.get('category')
                                    or sample.metadata.get('id') or sample.id or '')
            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)
    preds = list(await asyncio.gather(*(one(s) for s in work)))
    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,
    limit_per_task: Optional[int] = None,
    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.loop import ENV_REGISTRY, get_env

        if env not in ENV_REGISTRY:
            raise KeyError(f'unknown env {env!r}; available: {", ".join(ENV_REGISTRY.names())}')
        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,
            limit_per_task=limit_per_task,
            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
    return report


def _make_adapter(spec: str) -> ModelAdapter:
    """'mock:boxed' -> MockAdapter(mode='boxed'); else resolve_adapter().

    The colon-mode syntax exists ONLY for 'mock': adapter names contain no
    scheme/colon, so 'mock:xxx' is safe while URLs ('openai/http://...')
    must never be split on ':'.
    """
    base, sep, mode = spec.partition(':')
    if sep and '/' not in base and base == 'mock':
        adapter = resolve_adapter('mock')
        adapter.extra['mode'] = mode or 'echo'
        return adapter
    return resolve_adapter(spec)


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