import argparse
import os
import os as _os_env_toggle
import random

import numpy as np
import soundfile as sf
import torch
from utils import make_omni_prompt
from vllm.sampling_params import SamplingParams

from vllm_omni.entrypoints.omni_llm import OmniLLM

_os_env_toggle.environ["VLLM_USE_V1"] = "1"

SEED = 42
# Set all random seeds
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

# Make PyTorch deterministic
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# Set environment variables for deterministic behavior
os.environ["PYTHONHASHSEED"] = str(SEED)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model",
        required=True,
        help="Path to merged model directory (will be created if downloading).",
    )
    parser.add_argument("--thinker-model", type=str, default=None)
    parser.add_argument("--talker-model", type=str, default=None)
    parser.add_argument("--code2wav-model", type=str, default=None)
    parser.add_argument(
        "--hf-hub-id",
        default="Qwen/Qwen2.5-Omni-7B",
        help="Hugging Face repo id to download if needed.",
    )
    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
    parser.add_argument(
        "--code2wav-dir",
        default=None,
        help="Path to code2wav folder (contains spk_dict.pt).",
    )
    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
    parser.add_argument("--max-model-len", type=int, default=32768)
    parser.add_argument(
        "--init-sleep-seconds",
        type=int,
        default=20,
        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
    )

    parser.add_argument("--thinker-only", action="store_true")
    parser.add_argument("--text-only", action="store_true")
    parser.add_argument("--do-wave", action="store_true")
    parser.add_argument(
        "--prompt_type",
        choices=[
            "text",
            "audio",
            "audio-long",
            "audio-long-chunks",
            "audio-long-expand-chunks",
            "image",
            "video",
            "video-frames",
            "audio-in-video",
            "audio-in-video-v2",
            "audio-multi-round",
            "badcase-vl",
            "badcase-text",
            "badcase-image-early-stop",
            "badcase-two-audios",
            "badcase-two-videos",
            "badcase-multi-round",
            "badcase-voice-type",
            "badcase-voice-type-v2",
            "badcase-audio-tower-1",
            "badcase-audio-only",
        ],
        default="text",
    )
    parser.add_argument("--use-torchvision", action="store_true")
    parser.add_argument("--tokenize", action="store_true")
    parser.add_argument(
        "--output-wav",
        default="output.wav",
        help="[Deprecated] Output wav directory (use --output-dir).",
    )
    parser.add_argument(
        "--output-dir",
        default="outputs",
        help="Output directory to save text and wav files together.",
    )
    parser.add_argument(
        "--thinker-hidden-states-dir",
        default="thinker_hidden_states",
        help="Path to thinker hidden states directory.",
    )
    parser.add_argument(
        "--batch-timeout",
        type=int,
        default=5,
        help="Timeout for batching in seconds (default: 5)",
    )
    parser.add_argument(
        "--init-timeout",
        type=int,
        default=300,
        help="Timeout for initializing stages in seconds (default: 300)",
    )
    parser.add_argument(
        "--shm-threshold-bytes",
        type=int,
        default=65536,
        help="Threshold for using shared memory in bytes (default: 65536)",
    )
    parser.add_argument(
        "--enable-stats",
        action="store_true",
        default=False,
        help="Enable writing detailed statistics (default: disabled)",
    )
    parser.add_argument(
        "--txt-prompts",
        type=str,
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
    parser.add_argument(
        "--worker-backend",
        type=str,
        default="process",
        choices=["process","ray"],
        help="backend"
    )
    parser.add_argument(
        "--ray-address",
        type=str,
        default=None,
        help="Path to a .txt file with one prompt per line (preferred)."
        )

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    model_name = args.model
    try:
        # Preferred: load from txt file (one prompt per line)
        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
            with open(args.txt_prompts, encoding="utf-8") as f:
                lines = [ln.strip() for ln in f.readlines()]
            args.prompts = [ln for ln in lines if ln != ""]
            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
    except Exception as e:
        print(f"[Error] Failed to load prompts: {e}")
        raise

    if args.prompts is None:
        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
    omni_llm = OmniLLM(
        model=model_name,
        log_stats=args.enable_stats,
        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
        init_sleep_seconds=args.init_sleep_seconds,
        batch_timeout=args.batch_timeout,
        init_timeout=args.init_timeout,
        shm_threshold_bytes=args.shm_threshold_bytes,
        worker_backend=args.worker_backend,
        ray_address=args.ray_address
    )
    thinker_sampling_params = SamplingParams(
        temperature=0.0,  # Deterministic - no randomness
        top_p=1.0,  # Disable nucleus sampling
        top_k=-1,  # Disable top-k sampling
        max_tokens=2048,
        seed=SEED,  # Fixed seed for sampling
        detokenize=True,
        repetition_penalty=1.1,
    )
    talker_sampling_params = SamplingParams(
        temperature=0.9,
        top_p=0.8,
        top_k=40,
        max_tokens=2048,
        seed=SEED,  # Fixed seed for sampling
        detokenize=True,
        repetition_penalty=1.05,
        stop_token_ids=[8294],
    )
    code2wav_sampling_params = SamplingParams(
        temperature=0.0,  # Deterministic - no randomness
        top_p=1.0,  # Disable nucleus sampling
        top_k=-1,  # Disable top-k sampling
        max_tokens=2048,
        seed=SEED,  # Fixed seed for sampling
        detokenize=True,
        repetition_penalty=1.1,
    )

    sampling_params_list = [
        thinker_sampling_params,
        talker_sampling_params,
        code2wav_sampling_params,
    ]
    import time
    for i in range(1):
        t1 = time.time()
        prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
        print(f"prompt:{prompt}")
        omni_outputs = omni_llm.generate(prompt, sampling_params_list)
        t2 = time.time()
        print(f"==========> time:{t2-t1}")

    # Determine output directory: prefer --output-dir; fallback to --output-wav
    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
    os.makedirs(output_dir, exist_ok=True)
    for stage_outputs in omni_outputs:
        if stage_outputs.final_output_type == "text":
            for output in stage_outputs.request_output:
                request_id = int(output.request_id)
                text_output = output.outputs[0].text
                # Save aligned text file per request
                prompt_text = args.prompts[request_id]
                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
                lines = []
                lines.append("Prompt:\n")
                lines.append(str(prompt_text) + "\n")
                lines.append("vllm_text_output:\n")
                lines.append(str(text_output).strip() + "\n")
                try:
                    with open(out_txt, "w", encoding="utf-8") as f:
                        f.writelines(lines)
                except Exception as e:
                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
                print(f"Request ID: {request_id}, Text saved to {out_txt}")
        elif stage_outputs.final_output_type == "audio":
            for output in stage_outputs.request_output:
                request_id = int(output.request_id)
                audio_tensor = output.multimodal_output["audio"]
                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
                print(f"Request ID: {request_id}, Saved audio to {output_wav}")


if __name__ == "__main__":
    main()