"""
Batch inference for TTS evaluation benchmarks (Seed-TTS-eval & CV3-eval).

Uses the llama.cpp backend (MOSS-TTS-Delay).

Expected benchmark layout (per case)::

    {benchmark_dir}/{task}/{case_id}/prompt.wav
    {benchmark_dir}/{task}/{case_id}/label.txt

Output layout::

    {result_dir}/{task}/{case_id}/pred.wav

Usage::

    python scripts/batch_eval_llama_cpp.py \\
        --config configs/llama_cpp/default.yaml \\
        --benchmark-dir /path/to/eval/tts \\
        --result-dir results/my_run \\
        --tasks seed-tts-zeroshot-zh seed-tts-zeroshot-en
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import soundfile as sf
from tqdm import tqdm

from moss_tts_delay.llama_cpp import LlamaCppPipeline, PipelineConfig
from moss_tts_delay.llama_cpp._constants import SAMPLE_RATE

log = logging.getLogger(__name__)

SEED_TTS_TASKS = [
    "seed-tts-zeroshot-zh",
    "seed-tts-zeroshot-en",
    "seed-tts-zeroshot-hard-zh",
]

CV3_TASKS = [
    "cv3-crosslingual-en",
    "cv3-crosslingual-hard-en",
    "cv3-zeroshot-en",
    "cv3-zeroshot-hard-en",
    "cv3-crosslingual-zh",
    "cv3-crosslingual-hard-zh",
    "cv3-zeroshot-zh",
    "cv3-zeroshot-hard-zh",
]

ALL_TASKS = SEED_TTS_TASKS + CV3_TASKS + ["demo-zh", "demo-en"]

TASK_LANGUAGE = {
    "seed-tts-zeroshot-zh": "zh",
    "seed-tts-zeroshot-en": "en",
    "seed-tts-zeroshot-hard-zh": "zh",
    "cv3-crosslingual-en": "en",
    "cv3-crosslingual-hard-en": "en",
    "cv3-zeroshot-en": "en",
    "cv3-zeroshot-hard-en": "en",
    "cv3-crosslingual-zh": "zh",
    "cv3-crosslingual-hard-zh": "zh",
    "cv3-zeroshot-zh": "zh",
    "cv3-zeroshot-hard-zh": "zh",
    "demo-zh": "zh",
    "demo-en": "en",
}


@dataclass
class CaseResult:
    task: str
    case_id: str
    success: bool
    audio_duration: float = 0.0
    generation_time: float = 0.0
    error: str = ""


def discover_cases(benchmark_dir: Path, tasks: list[str]) -> list[tuple[str, str, Path, str]]:
    cases = []
    for task in tasks:
        task_dir = benchmark_dir / task
        if not task_dir.is_dir():
            log.warning("Task directory not found: %s", task_dir)
            continue
        for case_dir in sorted(task_dir.iterdir()):
            if not case_dir.is_dir():
                continue
            prompt_wav = case_dir / "prompt.wav"
            label_txt = case_dir / "label.txt"
            if not label_txt.exists():
                log.warning("Missing label.txt: %s", case_dir)
                continue
            text = label_txt.read_text().strip()
            cases.append((task, case_dir.name, prompt_wav, text))
    return cases


def run_batch(
    pipeline: LlamaCppPipeline,
    cases: list[tuple[str, str, Path, str]],
    result_dir: Path,
    max_cases: int = 0,
    skip_existing: bool = True,
) -> list[CaseResult]:
    results: list[CaseResult] = []
    total = len(cases) if max_cases <= 0 else min(max_cases, len(cases))
    cases = cases[:total]

    log.info("Running %d evaluation cases, output -> %s", total, result_dir)

    pbar = tqdm(cases, desc="Evaluation", unit="case", total=total, dynamic_ncols=True)
    for i, (task, case_id, prompt_wav, text) in enumerate(pbar):
        pbar.set_postfix_str(f"{task}/{case_id}")
        out_dir = result_dir / task / case_id
        out_wav = out_dir / "pred.wav"

        if skip_existing and out_wav.exists():
            log.info("[%d/%d] %s/%s — skipped (exists)", i + 1, total, task, case_id)
            results.append(CaseResult(task=task, case_id=case_id, success=True))
            continue

        log.info("[%d/%d] %s/%s — %s", i + 1, total, task, case_id, text[:60])
        t0 = time.time()

        try:
            lang = TASK_LANGUAGE.get(task)
            ref_audio = str(prompt_wav) if prompt_wav.exists() else None

            waveform = pipeline.generate(
                text=text, reference_audio=ref_audio, language=lang,
            )
            elapsed = time.time() - t0

            if waveform.size == 0:
                results.append(CaseResult(
                    task=task, case_id=case_id, success=False,
                    generation_time=elapsed, error="empty waveform",
                ))
                continue

            out_dir.mkdir(parents=True, exist_ok=True)
            sf.write(str(out_wav), waveform, SAMPLE_RATE)
            audio_dur = len(waveform) / SAMPLE_RATE

            results.append(CaseResult(
                task=task, case_id=case_id, success=True,
                audio_duration=audio_dur, generation_time=elapsed,
            ))
            log.info(
                "  -> %.2fs audio in %.2fs (RTF=%.2f)",
                audio_dur, elapsed, elapsed / max(audio_dur, 1e-6),
            )

        except Exception as e:
            elapsed = time.time() - t0
            log.error("  -> FAILED: %s", e)
            results.append(CaseResult(
                task=task, case_id=case_id, success=False,
                generation_time=elapsed, error=str(e),
            ))

    return results


def write_summary(results: list[CaseResult], result_dir: Path) -> None:
    succeeded = [r for r in results if r.success]
    failed = [r for r in results if not r.success]

    per_task: dict[str, dict] = {}
    for r in results:
        if r.task not in per_task:
            per_task[r.task] = {"total": 0, "success": 0, "failed": 0, "total_audio_s": 0.0, "total_gen_s": 0.0}
        per_task[r.task]["total"] += 1
        if r.success:
            per_task[r.task]["success"] += 1
            per_task[r.task]["total_audio_s"] += r.audio_duration
            per_task[r.task]["total_gen_s"] += r.generation_time
        else:
            per_task[r.task]["failed"] += 1

    for task, stats in per_task.items():
        if stats["total_audio_s"] > 0:
            stats["avg_rtf"] = round(stats["total_gen_s"] / stats["total_audio_s"], 3)

    summary = {
        "total_cases": len(results),
        "succeeded": len(succeeded),
        "failed": len(failed),
        "per_task": per_task,
    }

    if failed:
        summary["failures"] = [
            {"task": r.task, "case_id": r.case_id, "error": r.error}
            for r in failed
        ]

    summary_path = result_dir / "inference_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    log.info("Summary written to %s", summary_path)

    print("\n" + "=" * 60)
    print("  BATCH INFERENCE SUMMARY")
    print("=" * 60)
    print(f"  Total:     {len(results)}")
    print(f"  Succeeded: {len(succeeded)}")
    print(f"  Failed:    {len(failed)}")
    for task, stats in per_task.items():
        rtf = stats.get("avg_rtf", "N/A")
        print(f"  {task}: {stats['success']}/{stats['total']}  RTF={rtf}")
    print("=" * 60 + "\n")


def main():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    parser = argparse.ArgumentParser(
        description="Batch TTS evaluation (llama.cpp backend)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--config", required=True, help="Pipeline YAML config")
    parser.add_argument(
        "--benchmark-dir",
        default="/inspire/hdd/project/embodied-multimodality/public/speech_generation/data/eval/tts",
    )
    parser.add_argument("--result-dir", required=True)
    parser.add_argument("--tasks", nargs="+", default=None)
    parser.add_argument("--suite", choices=["seed-tts", "cv3", "all"], default=None)
    parser.add_argument("--max-cases", type=int, default=0)
    parser.add_argument("--no-skip", action="store_true")

    parser.add_argument("--text-temp", type=float, default=None)
    parser.add_argument("--audio-temp", type=float, default=None)
    parser.add_argument("--audio-top-p", type=float, default=None)
    parser.add_argument("--audio-top-k", type=int, default=None)
    parser.add_argument("--audio-rep-penalty", type=float, default=None)
    parser.add_argument("--n-gpu-layers", type=int, default=None)
    parser.add_argument("--max-tokens", type=int, default=None)
    parser.add_argument("--heads-backend", choices=["auto", "numpy", "torch"], default=None)

    args = parser.parse_args()
    config = PipelineConfig.from_yaml(args.config)

    if args.text_temp is not None:
        config.text_temperature = args.text_temp
    if args.audio_temp is not None:
        config.audio_temperature = args.audio_temp
    if args.audio_top_p is not None:
        config.audio_top_p = args.audio_top_p
    if args.audio_top_k is not None:
        config.audio_top_k = args.audio_top_k
    if args.audio_rep_penalty is not None:
        config.audio_repetition_penalty = args.audio_rep_penalty
    if args.n_gpu_layers is not None:
        config.n_gpu_layers = args.n_gpu_layers
    if args.max_tokens is not None:
        config.max_new_tokens = args.max_tokens
    if args.heads_backend is not None:
        config.heads_backend = args.heads_backend

    if args.tasks:
        tasks = args.tasks
    elif args.suite == "seed-tts":
        tasks = SEED_TTS_TASKS
    elif args.suite == "cv3":
        tasks = CV3_TASKS
    else:
        tasks = ALL_TASKS

    for t in tasks:
        if t not in ALL_TASKS:
            log.error("Unknown task: %s. Valid tasks: %s", t, ALL_TASKS)
            sys.exit(1)

    benchmark_dir = Path(args.benchmark_dir)
    result_dir = Path(args.result_dir)
    result_dir.mkdir(parents=True, exist_ok=True)

    cases = discover_cases(benchmark_dir, tasks)
    if not cases:
        log.error("No cases found in %s for tasks %s", benchmark_dir, tasks)
        sys.exit(1)
    log.info("Discovered %d cases across %d tasks", len(cases), len(tasks))

    run_meta = {
        "config": args.config,
        "benchmark_dir": str(benchmark_dir),
        "tasks": tasks,
        "sampling": {
            "text_temperature": config.text_temperature,
            "text_top_p": config.text_top_p,
            "text_top_k": config.text_top_k,
            "audio_temperature": config.audio_temperature,
            "audio_top_p": config.audio_top_p,
            "audio_top_k": config.audio_top_k,
            "audio_repetition_penalty": config.audio_repetition_penalty,
        },
        "max_new_tokens": config.max_new_tokens,
        "backbone_gguf": config.backbone_gguf,
        "heads_backend": config.heads_backend,
    }
    with open(result_dir / "run_meta.json", "w") as f:
        json.dump(run_meta, f, indent=2, ensure_ascii=False)

    with LlamaCppPipeline(config) as pipeline:
        results = run_batch(
            pipeline, cases, result_dir,
            max_cases=args.max_cases,
            skip_existing=not args.no_skip,
        )

    write_summary(results, result_dir)


if __name__ == "__main__":
    main()