#!/usr/bin/env python3
"""
Passo 4: Inference - Speech-to-Speech
Load trained model and generate audio responses from audio input.
"""

import os
import sys
import argparse
import torch
import torchaudio
import numpy as np
from pathlib import Path

# SNAC token configuration
SNAC_BASE = 128266
SNAC_MAX = SNAC_BASE + 7 * 4096  # 7 positions per frame, 4096 tokens each
EOS_TOKEN = 128009

def load_models(checkpoint_path: str, device: str = "cuda"):
    """Load all models for inference."""
    from transformers import WhisperModel, WhisperFeatureExtractor, AutoTokenizer
    import snac

    print("Loading models...")

    # Load Whisper encoder
    print("  Whisper encoder...")
    whisper_model = WhisperModel.from_pretrained("openai/whisper-large-v3").to(device)
    whisper_model.eval()
    feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-large-v3")

    # Load checkpoint
    print(f"  Loading checkpoint: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)

    # Debug: show checkpoint keys
    print(f"  Checkpoint keys: {list(checkpoint.keys())}")

    # Load LLM with LoRA
    print("  Loading LLM...")
    from transformers import AutoModelForCausalLM

    llm = AutoModelForCausalLM.from_pretrained(
        "canopylabs/3b-es_it-ft-research_release",
        torch_dtype=torch.bfloat16,
        device_map=device
    )

    # Load LoRA weights if present
    lora_loaded = False
    for lora_key in ['lora', 'lora_state_dict']:
        if lora_key in checkpoint:
            print(f"  Found LoRA weights with key '{lora_key}'")
            from peft import LoraConfig, get_peft_model
            lora_config = LoraConfig(
                r=16,
                lora_alpha=32,
                target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
                lora_dropout=0.0,  # No dropout for inference
                bias="none",
                task_type="CAUSAL_LM"
            )
            llm = get_peft_model(llm, lora_config)

            # Load the LoRA state dict
            lora_state = checkpoint[lora_key]
            print(f"  LoRA state dict has {len(lora_state)} keys")
            if len(lora_state) > 0:
                print(f"  Sample keys: {list(lora_state.keys())[:3]}")

            # Try to load with strict=False to see what matches
            result = llm.load_state_dict(lora_state, strict=False)
            print(f"  LoRA load - Missing: {len(result.missing_keys)}, Unexpected: {len(result.unexpected_keys)}")
            if result.missing_keys:
                print(f"    Missing keys (first 3): {result.missing_keys[:3]}")
            if result.unexpected_keys:
                print(f"    Unexpected keys (first 3): {result.unexpected_keys[:3]}")

            lora_loaded = True
            break

    if not lora_loaded:
        print("  WARNING: No LoRA weights found in checkpoint! Model will use base weights only.")
        print("  This might result in no SNAC token generation.")

    llm.eval()

    # Load Speech Adapter
    print("  Loading Speech Adapter...")
    from passo2_finetune_stage1 import SpeechAdapter
    adapter = SpeechAdapter(
        whisper_dim=1280,
        llm_dim=3072,
        downsample=5
    ).to(device)

    adapter_loaded = False
    for adapter_key in ['adapter', 'adapter_state_dict']:
        if adapter_key in checkpoint:
            print(f"  Found adapter weights with key '{adapter_key}'")
            result = adapter.load_state_dict(checkpoint[adapter_key], strict=False)
            print(f"  Adapter load - Missing: {len(result.missing_keys)}, Unexpected: {len(result.unexpected_keys)}")
            adapter_loaded = True
            break

    if not adapter_loaded:
        print("  WARNING: No adapter weights found in checkpoint!")

    adapter.eval()

    # Load SNAC decoder
    print("  SNAC decoder...")
    snac_model = snac.SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
    snac_model.eval()

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained("canopylabs/3b-es_it-ft-research_release")

    # Debug tokenizer info
    print(f"  Tokenizer vocab size: {tokenizer.vocab_size}")
    print(f"  BOS token: {tokenizer.bos_token} (id={tokenizer.bos_token_id})")
    print(f"  EOS token: {tokenizer.eos_token} (id={tokenizer.eos_token_id})")
    print(f"  SNAC token range: {SNAC_BASE} to {SNAC_BASE + 3*4096 - 1}")

    print("Models loaded!")
    return whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer


def encode_audio(audio_path: str, whisper_model, feature_extractor, adapter, device: str):
    """Encode input audio to embeddings."""
    # Load audio
    waveform, sr = torchaudio.load(audio_path)
    if sr != 16000:
        waveform = torchaudio.functional.resample(waveform, sr, 16000)

    # Extract Whisper features
    with torch.no_grad():
        inputs = feature_extractor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt")
        whisper_features = whisper_model.encoder(
            inputs.input_features.to(device)
        ).last_hidden_state

        # Adapt to LLM space
        adapted = adapter(whisper_features)

    return adapted.to(torch.bfloat16)


def decode_snac_tokens(tokens: list, snac_model, device: str):
    """Decode SNAC tokens to audio waveform.

    SNAC uses 3 hierarchical layers with 1:2:4 ratio.
    Each "frame" has 7 tokens in order:
      - 1 token from layer 0 (position 0)
      - 2 tokens from layer 1 (positions 1, 2)
      - 4 tokens from layer 2 (positions 3, 4, 5, 6)

    Tokens are offset by: SNAC_BASE + (position % 7) * 4096
    """
    if len(tokens) == 0:
        print("  Warning: No SNAC tokens to decode")
        return np.zeros(24000, dtype=np.float32)  # 1 second of silence

    layer0_tokens = []
    layer1_tokens = []
    layer2_tokens = []

    # Parse tokens by removing position-based offsets
    for i, tok in enumerate(tokens):
        pos = i % 7
        # Remove the offset to get the original code
        original = tok - SNAC_BASE - (pos * 4096)

        # Ensure codes are valid (0-4095)
        if original < 0 or original >= 4096:
            print(f"  Warning: Invalid code {original} at position {i} (token={tok}, pos={pos})")
            original = max(0, min(4095, original))

        if pos == 0:
            layer0_tokens.append(original)
        elif pos in [1, 2]:
            layer1_tokens.append(original)
        else:  # pos in [3, 4, 5, 6]
            layer2_tokens.append(original)

    # Calculate how many complete frames we have
    n_frames = len(tokens) // 7

    if n_frames == 0:
        print(f"  Warning: Not enough tokens for a complete frame ({len(tokens)} tokens)")
        return np.zeros(24000, dtype=np.float32)

    print(f"  Decoding {n_frames} frames ({len(layer0_tokens)} L0, {len(layer1_tokens)} L1, {len(layer2_tokens)} L2)")

    # Truncate to complete frames
    layer0_tokens = layer0_tokens[:n_frames]
    layer1_tokens = layer1_tokens[:n_frames * 2]
    layer2_tokens = layer2_tokens[:n_frames * 4]

    codes = [
        torch.tensor([layer0_tokens], dtype=torch.long, device=device),
        torch.tensor([layer1_tokens], dtype=torch.long, device=device),
        torch.tensor([layer2_tokens], dtype=torch.long, device=device)
    ]

    # Decode
    with torch.no_grad():
        audio = snac_model.decode(codes)

    return audio.squeeze().cpu().numpy()


def generate_response(
    audio_input: str,
    whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer,
    device: str,
    max_new_tokens: int = 500,
    temperature: float = 0.7,
    debug: bool = True,
    use_prompt: bool = False
):
    """Generate speech response from audio input."""

    # Encode input audio
    print("Encoding input audio...")
    audio_embeddings = encode_audio(audio_input, whisper_model, feature_extractor, adapter, device)
    print(f"  Audio embeddings shape: {audio_embeddings.shape}")

    # Generate with LLM
    print("Generating response...")

    with torch.no_grad():
        # Get embeddings layer
        embed_layer = llm.get_input_embeddings()

        if use_prompt:
            # Option A: Add BOS token after audio
            bos_id = tokenizer.bos_token_id or 128000
            prompt_tokens = torch.tensor([[bos_id]], dtype=torch.long, device=device)
            prompt_embeds = embed_layer(prompt_tokens)
            input_embeds = torch.cat([audio_embeddings.to(torch.bfloat16), prompt_embeds.to(torch.bfloat16)], dim=1)
        else:
            # Option B: Just audio embeddings (as trained)
            # During training, target directly follows audio embeddings
            input_embeds = audio_embeddings.to(torch.bfloat16)

        print(f"  Input embeds shape: {input_embeds.shape}")

        # Create attention mask
        attention_mask = torch.ones(input_embeds.shape[:2], dtype=torch.long, device=device)

        # Different generation strategies
        # Strategy 1: Pure greedy (most stable)
        if temperature <= 0.01:
            outputs = llm.generate(
                inputs_embeds=input_embeds,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=EOS_TOKEN,
                use_cache=True
            )
        # Strategy 2: Sampling with nucleus
        else:
            outputs = llm.generate(
                inputs_embeds=input_embeds,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True,
                top_p=0.9,
                top_k=50,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=EOS_TOKEN,
                use_cache=True
            )

    # Extract generated tokens
    generated_tokens = outputs[0].tolist()

    if debug:
        print(f"\n=== DEBUG: Generated {len(generated_tokens)} total tokens ===")
        # Show first 50 tokens
        print(f"First 50 tokens: {generated_tokens[:50]}")
        # Show last 20 tokens
        print(f"Last 20 tokens: {generated_tokens[-20:]}")

        # Show token ranges with correct SNAC range
        text_tokens = [t for t in generated_tokens if t < SNAC_BASE]
        snac_range_tokens = [t for t in generated_tokens if SNAC_BASE <= t < SNAC_MAX]
        other_high = [t for t in generated_tokens if t >= SNAC_MAX]

        print(f"\nToken distribution:")
        print(f"  Text tokens (<{SNAC_BASE}): {len(text_tokens)}")
        print(f"  SNAC tokens ({SNAC_BASE}-{SNAC_MAX}): {len(snac_range_tokens)}")
        print(f"  Other high tokens (>={SNAC_MAX}): {len(other_high)}")

        # Try to decode text tokens
        if text_tokens:
            try:
                decoded_text = tokenizer.decode(text_tokens, skip_special_tokens=False)
                print(f"\nDecoded text: {decoded_text[:500]}")
            except Exception as e:
                print(f"Could not decode text: {e}")

        # Show some SNAC tokens if any
        if snac_range_tokens:
            print(f"\nFirst 20 SNAC tokens: {snac_range_tokens[:20]}")

    # Extract SNAC tokens from output (correct range)
    snac_tokens = [t for t in generated_tokens if SNAC_BASE <= t < SNAC_MAX]

    print(f"\nGenerated {len(snac_tokens)} SNAC tokens")

    # Decode to audio
    print("Decoding to audio...")
    audio_output = decode_snac_tokens(snac_tokens, snac_model, device)

    return audio_output


def main():
    parser = argparse.ArgumentParser(description="Speech-to-Speech Inference")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint")
    parser.add_argument("--input", type=str, required=True, help="Input audio file")
    parser.add_argument("--output", type=str, default="output.wav", help="Output audio file")
    parser.add_argument("--max_tokens", type=int, default=1000, help="Max tokens to generate")
    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature (0 for greedy)")
    parser.add_argument("--use_prompt", action="store_true", help="Add BOS token after audio embeddings")
    parser.add_argument("--no_debug", action="store_true", help="Disable debug output")
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}")

    # Load models
    whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer = load_models(
        args.checkpoint, device
    )

    # Generate
    audio_output = generate_response(
        args.input,
        whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer,
        device,
        max_new_tokens=args.max_tokens,
        temperature=args.temperature,
        debug=not args.no_debug,
        use_prompt=args.use_prompt
    )

    # Save output
    import soundfile as sf
    sf.write(args.output, audio_output, 24000)
    print(f"Saved output to: {args.output}")


if __name__ == "__main__":
    main()