"""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,
) -> tuple:
    """Fan out model calls; returns (raws, total_usage).

    Each sample becomes one user message (or its ChatMessage list is used
    verbatim for multi-turn samples). Tool declarations from sample.tools
    are passed through so fc/agent recipes degrade gracefully today and
    agent loops can reuse this adapter untouched.
    """
    gen_kwargs = gen_kwargs or {}
    sem = asyncio.Semaphore(concurrency)
    total_usage = Usage()
    done_count = 0
    t0 = time.time()
    raws: List[str] = []
    usages: List[Dict[str, Any]] = []

    async def one(sample: Sample) -> tuple:
        nonlocal done_count, total_usage
        messages = ([ChatMessage(role='user', content=sample.input)] 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') == 'boxed' \
                and sample.target not in ('', None):
            # oracle channel for mock:boxed so full pipelines verify 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)
        usage = out.usage.model_dump()
        done_count += 1
        if progress and (done_count % 20 == 0 or done_count == len(samples)):
            rate = done_count / max(time.time() - t0, 1e-6)
            print(f'  [{done_count}/{len(samples)}] {rate:.1f} samples/s '
                  f'tokens={total_usage.total_tokens}', flush=True)
        return text, usage

    work = samples[:limit] if limit else samples
    pairs = await asyncio.gather(*(one(s) for s in work))
    raws = [p[0] for p in pairs]
    usages = [p[1] for p in pairs]
    return raws, usages, total_usage


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,
) -> EvalReport:
    """Generate + score in one call. Model spec examples:
    'mock', 'mock:boxed', 'openai/http://gpu03:8000/v1?qwen3-8b', 'deploy:vllm/qwen3-8b'.
    """
    adapter = _make_adapter(model_spec)
    spec = getattr(dataset, 'spec', None)
    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)
    if progress:
        print(f'generating: {adapter} on {len(samples)} samples '
              f'(concurrency={concurrency})', flush=True)
    try:
        raws, usages, usage = await generate_predictions(adapter, samples, concurrency,
                                                         progress=progress, gen_kwargs=gen_kwargs)
    finally:
        await adapter.close()
    if judge is None and judge_spec:
        judge_adapter = _make_adapter(judge_spec)
        judge = _judge_callable(judge_adapter)

    preds = [{'raw': r, 'usage': u} for r, u in zip(raws, usages)]
    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):
    async def ask(messages) -> str:
        out = await judge_adapter.generate([ChatMessage(role='user', content=str(m)) for m in messages]
                                           if isinstance(messages, list) and messages and isinstance(messages[0], dict)
                                           else messages)
        return out.text

    import asyncio

    def sync_ask(messages):
        return asyncio.run(ask(messages))

    return sync_ask