#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Generic Thinker inference for speech benchmarks.

Input JSONL fields:
  id             required, stable sample id
  audio_path     required, absolute or cwd-relative wav path
  language       optional, en/zh/multi; defaults to en
  user_prompt    optional, task-specific text prompt/question
  system_prompt  optional, extra benchmark instruction
  benchmark      optional metadata
  task           optional metadata

Output JSONL fields are compatible with qwen3omni_thinker/talker_batch_abinfer.py:
  id, language, ability, audio_path, thinker_style, thinker_text, struct_pass
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from pathlib import Path
from typing import Any

import torch
from peft import PeftModel
from tqdm import tqdm
from transformers import Qwen3OmniMoeProcessor


DEFAULT_FRAMEWORK_ROOT = "/workspace/echoloc/TTS-Framework"
DEFAULT_BASE = "/workspace/echoloc/Qwen3-Omni-Instruct"
DEFAULT_LORA = (
    "/workspace/echoloc/Dataset/Novel/query_data/supervisor_iterations/iter_006/"
    "sft/thinker_emochange/iter006/markerw/output/latest_final_model"
)

PARSE_RE = re.compile(r"Style:\s*(.*?)\s*Text:\s*(.*)", re.S | re.I)

BASE_SYSTEM_PROMPT = """You are an emotionally intelligent spoken-dialogue model.
Listen to the user's input audio and follow the benchmark task instruction.

You must output exactly two fields:
Style: <voice style for the spoken answer>

Text: <the answer text that should be spoken>

Rules:
- Put the actual benchmark answer only in Text.
- If the task asks for transcription, Text must be the transcript only.
- If the task asks for a multiple-choice answer, Text must preserve the requested answer format, such as "The answer is: A".
- If the task asks for open-domain conversation, Text should be concise, helpful, and directly responsive.
- For empathy or role-play tasks, respond with warmth and emotional appropriateness while staying grounded in the user's audio.
- Style may describe speaker, emotion, pacing, and prosody, but must not contain the answer itself.
- Do not add Markdown, JSON, explanations, or extra fields."""


def load_rows(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def parse_style_text(text: str) -> tuple[str | None, str | None]:
    if not text:
        return None, None
    m = PARSE_RE.search(text)
    if not m:
        return None, None
    return m.group(1).strip(), m.group(2).strip()


def language_rule(language: str) -> str:
    if language == "zh":
        return "The input is Chinese. Text must answer in Chinese unless the task explicitly asks otherwise."
    if language == "multi":
        return "Preserve the language requested by the input audio or task."
    return "The input is English. Text must answer in English unless the task explicitly asks otherwise."


def build_conversation(item: dict[str, Any]) -> list[dict[str, Any]]:
    audio_path = str(item["audio_path"])
    language = str(item.get("language") or "en").lower()
    extra_system = (item.get("system_prompt") or "").strip()
    user_prompt = (item.get("user_prompt") or item.get("task_prompt") or "").strip()
    system_prompt = BASE_SYSTEM_PROMPT + "\n\n" + language_rule(language)
    if extra_system:
        system_prompt += "\n\nBenchmark instruction:\n" + extra_system
    user_text = user_prompt or "Please listen to the audio and respond directly to the user's request."
    return [
        {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
        {
            "role": "user",
            "content": [
                {"type": "audio", "audio": audio_path},
                {"type": "text", "text": user_text},
            ],
        },
    ]


def build_batch_inputs(processor: Qwen3OmniMoeProcessor, items: list[dict[str, Any]]):
    convs = [build_conversation(it) for it in items]
    processor.tokenizer.padding_side = "left"
    try:
        return processor.apply_chat_template(
            convs,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            padding=True,
        )
    except Exception as exc:
        print(f"[warn] batched apply_chat_template failed: {exc}; fallback to processor(audio=...).")

    import librosa

    texts: list[str] = []
    wavs: list[Any] = []
    for conv in convs:
        ap = None
        for content in conv[-1]["content"]:
            if content.get("type") == "audio":
                ap = content["audio"]
                break
        wav, _ = librosa.load(ap, sr=16000)
        wavs.append(wav)
        texts.append(processor.apply_chat_template(conv, add_generation_prompt=True, tokenize=False))
    return processor(text=texts, audio=wavs, sampling_rate=16000, return_tensors="pt", padding=True)


def validate_rows(rows: list[dict[str, Any]], no_check: bool = False) -> None:
    if not rows:
        raise ValueError("input JSONL has no rows")
    missing = []
    for idx, row in enumerate(rows[:20]):
        ap = row.get("audio_path")
        if not ap:
            missing.append(f"row {idx}: missing audio_path")
        elif not no_check and not os.path.isfile(str(ap)):
            missing.append(f"row {idx}: audio not found: {ap}")
    if missing:
        raise FileNotFoundError("\n".join(missing[:10]))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--input_jsonl", type=Path, required=True)
    ap.add_argument("--out_dir", type=Path, required=True)
    ap.add_argument("--framework_root", default=DEFAULT_FRAMEWORK_ROOT)
    ap.add_argument("--base", default=DEFAULT_BASE)
    ap.add_argument("--lora", default=DEFAULT_LORA)
    ap.add_argument("--base_only", action="store_true")
    ap.add_argument("--batch_size", type=int, default=8)
    ap.add_argument("--max_new_tokens", type=int, default=256)
    ap.add_argument("--num_shards", type=int, default=1)
    ap.add_argument("--shard_index", type=int, default=0)
    ap.add_argument("--smoke", type=int, default=0)
    ap.add_argument("--no_check", action="store_true")
    args = ap.parse_args()

    sys.path.insert(0, args.framework_root)
    from models.qwen3omni.model import Qwen3OmniThinkerModel  # noqa: WPS433

    rows = load_rows(args.input_jsonl)
    if args.smoke > 0:
        rows = rows[: args.smoke]
    if args.num_shards < 1:
        raise ValueError("--num_shards must be >= 1")
    if not (0 <= args.shard_index < args.num_shards):
        raise ValueError("--shard_index must satisfy 0 <= shard_index < num_shards")
    if args.num_shards > 1:
        before = len(rows)
        rows = [it for idx, it in enumerate(rows) if idx % args.num_shards == args.shard_index]
        print(f"[shard] selected {len(rows)}/{before} rows for {args.shard_index}/{args.num_shards}")
    validate_rows(rows, no_check=args.no_check)

    args.out_dir.mkdir(parents=True, exist_ok=True)
    out_jsonl = args.out_dir / "thinker_outputs.jsonl"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"[thinker] rows={len(rows)} device={device}")
    print(f"[thinker] base={args.base}")
    print(f"[thinker] lora={'BASE_ONLY' if args.base_only else args.lora}")
    processor = Qwen3OmniMoeProcessor.from_pretrained(args.base, trust_remote_code=True)
    base_wrapper = Qwen3OmniThinkerModel(
        model_path=args.base,
        torch_dtype=torch.bfloat16,
        enable_thinking=False,
    )
    base = base_wrapper.model.to(device)
    pad_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id
    eos_id = processor.tokenizer.eos_token_id
    if args.base_only:
        model = base
    else:
        model = PeftModel.from_pretrained(base, args.lora).merge_and_unload()
    model.eval()

    n_struct_ok = 0
    with out_jsonl.open("w", encoding="utf-8") as fout:
        pbar = tqdm(total=len(rows), desc="generic thinker")
        for start in range(0, len(rows), args.batch_size):
            batch = rows[start : start + args.batch_size]
            inputs = build_batch_inputs(processor, batch).to(device)
            if "input_features" in inputs:
                inputs["input_features"] = inputs["input_features"].to(torch.bfloat16)
            with torch.no_grad():
                gen = model.generate(
                    **inputs,
                    max_new_tokens=args.max_new_tokens,
                    do_sample=False,
                    pad_token_id=pad_id,
                    eos_token_id=eos_id,
                )
            in_len = inputs["input_ids"].shape[1]
            preds = processor.tokenizer.batch_decode(gen[:, in_len:], skip_special_tokens=True)
            for item, pred in zip(batch, preds):
                pred = pred.strip()
                style, text = parse_style_text(pred)
                struct_pass = style is not None and text is not None
                n_struct_ok += int(struct_pass)
                fout.write(
                    json.dumps(
                        {
                            **item,
                            "id": str(item.get("id")),
                            "language": item.get("language", "en"),
                            "ability": item.get("ability") or item.get("task") or item.get("benchmark", "benchmark"),
                            "thinker_style": style if struct_pass else "A clear, natural speaking voice; neutral and direct.",
                            "thinker_text": text if struct_pass else pred,
                            "struct_pass": struct_pass,
                            "raw_prediction": pred,
                        },
                        ensure_ascii=False,
                    )
                    + "\n"
                )
                fout.flush()
            pbar.update(len(batch))
        pbar.close()

    print(f"[thinker] done total={len(rows)} struct_pass={n_struct_ok} ({n_struct_ok / len(rows) * 100:.2f}%)")
    print(f"[thinker] output={out_jsonl}")


if __name__ == "__main__":
    main()