# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Single-sample MiniCPM-o native-duplex generation."""

from __future__ import annotations

import asyncio
import hashlib
import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

import pybase64 as base64

from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
from vllm_omni.experimental.fullduplex.client import (
    RealtimeDuplexClient,
    RealtimeEventCollector,
    build_realtime_url,
    wait_for,
)

if TYPE_CHECKING:
    from vllm_omni.clients.duplex import EventCollector

from .omni_duplex_eval_clock import extract_timed_sentences
from .omni_duplex_eval_media import iter_av_units, iter_jpegs, materialize_media, read_audio_pcm16, video_duration


@dataclass(frozen=True)
class GenerateSampleResult:
    """Timed-sentence path plus duplex metrics for one generate sample."""

    output: Path
    request_metrics: list[dict[str, object]] = field(default_factory=list)
    session_metrics: dict[str, object] = field(default_factory=dict)


def _ref_audio(path: str | Path) -> str:
    value = Path(path).expanduser().read_bytes()
    return "data:audio/wav;base64," + base64.b64encode(value).decode("ascii")


def _event_collector_from_realtime(source: RealtimeEventCollector) -> EventCollector:
    """Replay experimental collector events onto the public EventCollector."""
    from vllm_omni.clients.duplex import EventCollector

    collector = EventCollector()
    for event, received_at_s in zip(source.events, source.event_received_at_s, strict=True):
        collector.add(event, received_at_s=received_at_s)
    return collector


def _with_sample_identity(payload: dict[str, object], sample: DuplexSample) -> dict[str, object]:
    return {"sample_id": sample.id, "split": sample.split, **payload}


async def generate_sample(
    sample: DuplexSample,
    *,
    url: str,
    model: str,
    ref_audio: str | Path,
    output_root: str | Path,
    fps: float = 1.0,
    mix: str = "question",
    pace: str = "realtime",
    clock: str = "media",
    overwrite: bool = False,
    unit_ms: int = 1000,
) -> GenerateSampleResult:
    """Generate one sample's timed sentences and collect duplex session metrics.

    Timed sentences and ``*.meta.json`` stay the quality-eval artifacts. Duplex
    request/session metrics are returned for the CLI to write
    ``duplex_metrics.json``; they are not written into meta.
    """
    output = Path(output_root) / sample.split / f"{sample.id}.json"
    meta_path = output.with_name(output.stem + ".meta.json")
    if output.exists() and not overwrite:
        return GenerateSampleResult(output=output)
    if mix != "question":
        raise NotImplementedError("v1 supports mix=question; soundtrack mixing is reserved for P1")
    media_dir = output.parent / ".media"
    audio_path = materialize_media(sample.question_audio, media_dir, f"{sample.id}_question", ".wav")
    video_path = materialize_media(sample.video, media_dir, sample.id, ".mp4")
    pcm = read_audio_pcm16(audio_path)
    duration = sample.video_duration or video_duration(video_path)
    frames = iter_jpegs(video_path, fps=fps, duration=duration)
    realtime = pace == "realtime"
    if pace not in {"realtime", "as-fast-as-possible"}:
        raise ValueError("pace must be realtime or as-fast-as-possible")
    client = RealtimeDuplexClient(build_realtime_url(url, model, autostart=False))
    response_done = False
    drain_timeout = None
    close_timeout = None
    stream_start: float | None = None
    async with client:
        await client.configure(model, ref_audio=_ref_audio(ref_audio), instructions="Streaming Omni Conversation.")
        ack_task = asyncio.create_task(_ack_playback(client))
        try:
            stream_start = time.monotonic()
            await client.stream_av_units(iter_av_units(pcm, frames, unit_ms=unit_ms), realtime=realtime)
            await client.commit()
            try:
                await wait_for(
                    lambda: client.events.count("response.done") > 0,
                    timeout_s=max(20.0, duration + 20.0),
                    label="response.done",
                )
                response_done = True
            except TimeoutError as exc:
                drain_timeout = str(exc)
        finally:
            ack_task.cancel()
            try:
                await ack_task
            except asyncio.CancelledError:
                pass
        await client.acknowledge_playback()
        try:
            await client.close_session(timeout_s=20.0)
        except TimeoutError as exc:
            close_timeout = str(exc)
        events = list(client.events.events)
        collector = _event_collector_from_realtime(client.events)
    timed = [sentence.as_dict() for sentence in extract_timed_sentences(events, clock=clock)]
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(timed, ensure_ascii=False, indent=2), encoding="utf-8")
    meta = {
        "id": sample.id,
        "split": sample.split,
        "clock": clock if realtime else "invalid",
        "pace": pace,
        "mix": mix,
        "fps": fps,
        "unit_ms": unit_ms,
        "model": model,
        "response_done": response_done,
        "drain_timeout": drain_timeout,
        "close_timeout": close_timeout,
        "ref_audio_sha256": hashlib.sha256(Path(ref_audio).read_bytes()).hexdigest(),
    }
    meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
    from vllm_omni.benchmarks.duplex_session_metrics import collect_duplex_session_metrics

    bundle = collect_duplex_session_metrics(
        collector,
        stream_start=stream_start if stream_start is not None else 0.0,
        session_id=sample.id,
    )
    return GenerateSampleResult(
        output=output,
        request_metrics=[_with_sample_identity(metric, sample) for metric in bundle.request_metrics],
        session_metrics=_with_sample_identity(bundle.session_metrics, sample),
    )


async def _ack_playback(client: RealtimeDuplexClient) -> None:
    while True:
        await asyncio.sleep(0.25)
        await client.acknowledge_playback()