#!/usr/bin/env python3
"""
Inference script for Speech-to-Speech model.

Usage:
    python inference.py --checkpoint ./checkpoints/stage2_best.pt --input audio.wav --output response.wav
    python inference.py --checkpoint ./checkpoints/stage2_best.pt --text "Hello, how are you?"
"""

import os
import sys
import argparse
import torch
import torch.nn as nn
import numpy as np

# SNAC token offsets for Orpheus
SNAC_BASE_OFFSET = 128266
EOS_TOKEN = 128009


class SpeechAdapter(nn.Module):
    """Same architecture as training - must match exactly."""
    def __init__(self, whisper_dim=1280, llm_dim=3072, downsample=5, intermediate_dim=2048):
        super().__init__()
        self.downsample = downsample
        concat_dim = whisper_dim * downsample

        self.ffn = nn.Sequential(
            nn.Linear(concat_dim, intermediate_dim),
            nn.GELU(),
            nn.Linear(intermediate_dim, llm_dim),
            nn.LayerNorm(llm_dim)
        )

    def forward(self, x):
        B, T, D = x.shape
        T_new = (T // self.downsample) * self.downsample
        x = x[:, :T_new]
        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
        return self.ffn(x)


def decode_snac_tokens(snac_tokens, device="cuda"):
    """Decode SNAC tokens to audio waveform."""
    try:
        from snac import SNAC

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

        # Remove offsets from tokens
        raw_tokens = []
        for i, tok in enumerate(snac_tokens):
            pos = i % 7
            offset = SNAC_BASE_OFFSET + pos * 4096
            raw_tok = tok - offset
            if 0 <= raw_tok < 4096:
                raw_tokens.append(raw_tok)

        if len(raw_tokens) == 0:
            return None, 24000

        # Reshape to SNAC format: 7 tokens per frame
        num_frames = len(raw_tokens) // 7
        if num_frames == 0:
            return None, 24000

        raw_tokens = raw_tokens[:num_frames * 7]

        # SNAC expects [batch, layers, time] - 3 layers with different rates
        # Layer 0: 1 token/frame, Layer 1: 2 tokens/frame, Layer 2: 4 tokens/frame
        codes = []
        for frame_idx in range(num_frames):
            base = frame_idx * 7
            codes.append(raw_tokens[base:base+7])

        codes = torch.tensor(codes, device=device)

        # Reorganize into SNAC layer format
        layer0 = codes[:, 0:1].T  # [1, num_frames]
        layer1 = codes[:, 1:3].T.reshape(1, -1)  # [1, num_frames*2]
        layer2 = codes[:, 3:7].T.reshape(1, -1)  # [1, num_frames*4]

        with torch.no_grad():
            audio = snac.decode([layer0, layer1, layer2])

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

    except Exception as e:
        print(f"SNAC decode error: {e}")
        return None, 24000


def extract_whisper_features(audio_path, device="cuda"):
    """Extract Whisper encoder features from audio file."""
    try:
        from transformers import WhisperProcessor, WhisperModel
        import librosa

        # Load audio
        audio, sr = librosa.load(audio_path, sr=16000)

        # Load Whisper
        processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
        model = WhisperModel.from_pretrained("openai/whisper-large-v3").to(device)
        model.eval()

        # Process
        inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
        input_features = inputs.input_features.to(device)

        with torch.no_grad():
            encoder_outputs = model.encoder(input_features)
            features = encoder_outputs.last_hidden_state

        return features

    except Exception as e:
        print(f"Whisper feature extraction error: {e}")
        return None


def generate_response(model, adapter, tokenizer, audio_embeds, device, max_new_tokens=500):
    """Generate interleaved text+audio response."""

    # Get the base model for generation
    if hasattr(model, 'get_base_model'):
        base_model = model.get_base_model()
    else:
        base_model = model

    # Start generation from audio embeddings
    generated_tokens = []

    # Create initial input from audio embeddings
    current_embeds = audio_embeds

    with torch.no_grad():
        for step in range(max_new_tokens):
            # Forward pass
            outputs = model(inputs_embeds=current_embeds, use_cache=False)
            logits = outputs.logits

            # Get next token (greedy)
            next_token_logits = logits[:, -1, :]
            next_token = torch.argmax(next_token_logits, dim=-1)

            token_id = next_token.item()
            generated_tokens.append(token_id)

            # Check for EOS
            if token_id == EOS_TOKEN:
                break

            # Get embedding for next token
            if hasattr(base_model, 'model'):
                next_embed = base_model.model.embed_tokens(next_token.unsqueeze(0))
            else:
                next_embed = base_model.embed_tokens(next_token.unsqueeze(0))

            # Append to current embeddings
            current_embeds = torch.cat([current_embeds, next_embed], dim=1)

            # Truncate if too long (keep last 2048 tokens)
            if current_embeds.shape[1] > 2048:
                current_embeds = current_embeds[:, -2048:]

    return generated_tokens


def separate_tokens(generated_tokens):
    """Separate text and audio tokens from interleaved output."""
    text_tokens = []
    audio_tokens = []

    for tok in generated_tokens:
        if tok >= SNAC_BASE_OFFSET:
            audio_tokens.append(tok)
        elif tok != EOS_TOKEN:
            text_tokens.append(tok)

    return text_tokens, audio_tokens


def main():
    parser = argparse.ArgumentParser(description="Speech-to-Speech Inference")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint (stage1 or stage2)")
    parser.add_argument("--input", type=str, default=None, help="Input audio file")
    parser.add_argument("--text", type=str, default=None, help="Input text (for testing without audio)")
    parser.add_argument("--output", type=str, default="./output.wav", help="Output audio file")
    parser.add_argument("--model_path", type=str, default="canopylabs/3b-es_it-ft-research_release")
    parser.add_argument("--max_tokens", type=int, default=500)
    parser.add_argument("--device", type=str, default=None)
    args = parser.parse_args()

    # Determine device
    if args.device:
        device = torch.device(args.device)
    elif torch.cuda.is_available():
        device = torch.device("cuda")
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    print(f"Device: {device}")

    # Determine dtype
    torch_dtype = torch.bfloat16 if device.type == 'cuda' else torch.float32

    # Load tokenizer
    print(f"Loading tokenizer: {args.model_path}")
    from transformers import AutoTokenizer, AutoModelForCausalLM
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)

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

    has_lora = "lora" in ckpt
    print(f"Checkpoint type: {'Stage 2 (Adapter + LoRA)' if has_lora else 'Stage 1 (Adapter only)'}")

    # Load LLM
    print(f"Loading LLM: {args.model_path}")
    llm = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        torch_dtype=torch_dtype,
        attn_implementation="sdpa",
    ).to(device)

    # Apply LoRA if Stage 2
    if has_lora:
        from peft import LoraConfig, get_peft_model, TaskType
        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,
            bias="none",
            task_type=TaskType.CAUSAL_LM
        )
        llm = get_peft_model(llm, lora_config)
        llm.load_state_dict(ckpt["lora"], strict=False)
        print("LoRA weights loaded")

    llm.eval()

    # Load adapter
    print("Loading adapter...")
    adapter = SpeechAdapter(
        whisper_dim=1280,
        llm_dim=3072,
        downsample=5,
        intermediate_dim=2048
    ).to(device, dtype=torch_dtype)
    adapter.load_state_dict(ckpt["adapter"])
    adapter.eval()
    print("Adapter loaded")

    # Get input embeddings
    if args.input:
        print(f"Processing audio: {args.input}")
        whisper_features = extract_whisper_features(args.input, device)
        if whisper_features is None:
            print("Failed to extract Whisper features")
            return
        audio_embeds = adapter(whisper_features.to(torch_dtype))

    elif args.text:
        print(f"Processing text: {args.text}")
        # For text input, create dummy audio embeddings (zeros)
        # This is just for testing the generation pipeline
        dummy_features = torch.randn(1, 100, 1280, device=device, dtype=torch_dtype)
        audio_embeds = adapter(dummy_features)

        # Optionally prepend text tokens
        text_tokens = tokenizer.encode(args.text, add_special_tokens=False)
        print(f"Text tokens: {text_tokens[:10]}...")
    else:
        print("ERROR: Provide --input (audio file) or --text")
        return

    print(f"Audio embeddings shape: {audio_embeds.shape}")

    # Generate response
    print(f"Generating response (max {args.max_tokens} tokens)...")
    generated_tokens = generate_response(
        llm, adapter, tokenizer, audio_embeds, device,
        max_new_tokens=args.max_tokens
    )

    print(f"Generated {len(generated_tokens)} tokens")

    # Separate text and audio
    text_tokens, audio_tokens = separate_tokens(generated_tokens)
    print(f"Text tokens: {len(text_tokens)}, Audio tokens: {len(audio_tokens)}")

    # Decode text
    if text_tokens:
        decoded_text = tokenizer.decode(text_tokens, skip_special_tokens=True)
        print(f"\nGenerated text: {decoded_text}")

    # Decode audio
    if audio_tokens:
        print(f"\nDecoding {len(audio_tokens)} audio tokens...")
        audio, sr = decode_snac_tokens(audio_tokens, device)

        if audio is not None:
            import soundfile as sf
            sf.write(args.output, audio, sr)
            print(f"Audio saved: {args.output}")
        else:
            print("Failed to decode audio")
    else:
        print("No audio tokens generated")

    # Show raw tokens for debugging
    print(f"\nFirst 20 generated tokens: {generated_tokens[:20]}")


if __name__ == "__main__":
    main()