#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from scripts.qwen35_qwen3tts_voiceclone_smoke import generate_portuguese_answer, load_llm
from speech_bridge_gemma.mimi_overfit import resolve_device, set_seed


def parse_args() -> argparse.Namespace:
    root = Path(os.environ.get("S2S_ROOT", "/workspace/gemma1b-tts-integration"))
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", default=str(root))
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--llm-model", default="Qwen/Qwen3.5-0.8B")
    parser.add_argument("--llm-dtype", choices=["auto", "float32", "float16", "bfloat16"], default="float16")
    parser.add_argument("--text-sft-checkpoint", default="job_output/qwen3-text-sft-full-ptbr-20260606T022749Z/qwen_text_sft_layers.pt")
    parser.add_argument("--system", default="Voce e uma assistente brasileira. Responda sempre em portugues brasileiro, em uma frase curta, direta e natural.")
    parser.add_argument("--max-answer-tokens", type=int, default=48)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--seed", type=int, default=1337)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    os.chdir(Path(args.root).resolve())
    payload = json.loads(sys.stdin.read() or "{}")
    question = str(payload.get("question") or "").strip()
    context = str(payload.get("context") or "").strip()
    if not question:
        raise ValueError("question is required")
    set_seed(args.seed)
    device = resolve_device(args.device)
    tokenizer, llm, system, sft = load_llm(args, device)
    answer = generate_portuguese_answer(
        llm=llm,
        tokenizer=tokenizer,
        question=question,
        system=system,
        device=device,
        max_new_tokens=args.max_answer_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        context=context,
    )
    print(
        json.dumps(
            {
                "answer": answer,
                "system": system,
                "text_sft_loaded_layers": sft.get("text_sft_loaded_layers", 0),
            },
            ensure_ascii=False,
        ),
        flush=True,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())