import modal
import os
import subprocess
import time
import json
import tempfile
import logging
import sys
import warnings
import shutil
import base64
import requests
from typing import Optional, Dict, Any, Union
from datetime import datetime
from types import SimpleNamespace
from io import BytesIO
from pathlib import Path
from fastapi import HTTPException
from pydantic import BaseModel

# Suppress warnings for cleaner output
warnings.filterwarnings("ignore")

# Modal App Configuration
app = modal.App("multitalk-inference")

# Create a Volume for model weights
volume = modal.Volume.from_name("multitalk-weights", create_if_missing=True)
MODEL_DIR = Path("/models")

# HuggingFace secret for authentication
# huggingface_secret = modal.Secret.from_name(
#     "huggingface-secret", required_keys=["HF_TOKEN"]
# )
huggingface_secret = ''

# Base image with system dependencies
base_image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install(
        "ffmpeg",
        "libsndfile1",
        "git",
        "wget",
        "curl",
        "libgl1-mesa-glx",
        "libglib2.0-0",
        "libsm6",
        "libxext6",
        "libxrender-dev",
        "libgomp1",
    )
    .pip_install(
        "torch>=2.0.0",
        "torchvision",
        "torchaudio",
        "numpy",
        "soundfile",
        "librosa",
        "pyloudnorm",
        "einops",
        "transformers>=4.30.0",
        "huggingface-hub",
        "Pillow",
        "fastapi[standard]",
        "pydantic",
        "requests",
        "accelerate",
        "diffusers",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

# Download image with additional dependencies
download_image = base_image.pip_install(
    "huggingface_hub[hf_transfer]",
    "requests",
    "tqdm",
)

# Pydantic models for API
class MultiTalkRequest(BaseModel):
    image: str  # base64 encoded image
    first_audio: str  # base64 encoded audio
    prompt: str = "A smiling man and woman wearing headphones sit in front of microphones, appearing to host a podcast."
    second_audio: Optional[str] = None  # base64 encoded audio (optional)
    num_frames: int = 81
    sampling_steps: int = 40
    seed: Optional[int] = None
    turbo: bool = True

class MultiTalkResponse(BaseModel):
    video: str  # base64 encoded video
    metadata: Dict[str, Any]

class ErrorResponse(BaseModel):
    error: str
    type: str

# Helper functions
def loudness_norm(audio_array, sr=16000, lufs=-23):
    import pyloudnorm as pyln
    meter = pyln.Meter(sr)
    loudness = meter.integrated_loudness(audio_array)
    if abs(loudness) > 100:
        return audio_array
    normalized_audio = pyln.normalize.loudness(audio_array, loudness, lufs)
    return normalized_audio

def extract_audio_from_video(filename, sample_rate=16000):
    """Extract audio from video file with robust error handling"""
    import librosa
    raw_audio_path = f"{os.path.splitext(os.path.basename(filename))[0]}.wav"
    ffmpeg_command = [
        "ffmpeg",
        "-y",
        "-i",
        str(filename),
        "-vn",
        "-acodec",
        "pcm_s16le",
        "-ar",
        str(sample_rate),
        "-ac",
        "2",
        raw_audio_path,
    ]
    subprocess.run(ffmpeg_command, check=True, capture_output=True)
    human_speech_array, sr = librosa.load(raw_audio_path, sr=sample_rate)
    human_speech_array = loudness_norm(human_speech_array, sr)
    os.remove(raw_audio_path)
    return human_speech_array

def audio_prepare_single(audio_path, sample_rate=16000):
    import librosa
    ext = os.path.splitext(audio_path)[1].lower()
    if ext in [".mp4", ".mov", ".avi", ".mkv"]:
        human_speech_array = extract_audio_from_video(audio_path, sample_rate)
        return human_speech_array
    else:
        human_speech_array, sr = librosa.load(audio_path, sr=sample_rate)
        human_speech_array = loudness_norm(human_speech_array, sr)
        return human_speech_array

def audio_prepare_multi(left_path, right_path, audio_type, sample_rate=16000):
    import numpy as np
    human_speech_array1 = audio_prepare_single(left_path)
    human_speech_array2 = audio_prepare_single(right_path)

    if audio_type == "para":
        new_human_speech1 = human_speech_array1
        new_human_speech2 = human_speech_array2
    elif audio_type == "add":
        new_human_speech1 = np.concatenate([
            human_speech_array1[:human_speech_array1.shape[0]],
            np.zeros(human_speech_array2.shape[0]),
        ])
        new_human_speech2 = np.concatenate([
            np.zeros(human_speech_array1.shape[0]),
            human_speech_array2[:human_speech_array2.shape[0]],
        ])
    sum_human_speechs = new_human_speech1 + new_human_speech2
    return new_human_speech1, new_human_speech2, sum_human_speechs

def get_embedding(speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device="cpu"):
    """Extract audio embeddings optimized for GPU processing"""
    import torch
    import numpy as np
    from einops import rearrange
    
    audio_duration = len(speech_array) / sr
    video_length = audio_duration * 25  # Assume the video fps is 25

    # Extract audio features
    audio_feature = np.squeeze(
        wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values
    )
    audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
    audio_feature = audio_feature.unsqueeze(0)

    # Generate embeddings on appropriate device
    with torch.no_grad():
        embeddings = audio_encoder(
            audio_feature, seq_len=int(video_length), output_hidden_states=True
        )

    if len(embeddings) == 0:
        print("⚠️ Failed to extract audio embedding")
        return None

    audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
    audio_emb = rearrange(audio_emb, "b s d -> s b d")

    # Keep on CPU for compatibility with downstream processing
    audio_emb = audio_emb.cpu().detach()
    return audio_emb

def download_with_resume(url: str, dest: str, max_retries: int = 3, chunk_size: int = 8192 * 1024) -> None:
    """Download file with resume capability and robust error handling."""
    print(f"[!] Downloading from: {url}")
    print(f"[~] Destination: {dest}")

    # Create destination directory if it doesn't exist
    os.makedirs(os.path.dirname(dest), exist_ok=True)

    # Get initial file size if resuming
    initial_size = 0
    if os.path.exists(dest):
        initial_size = os.path.getsize(dest)
        print(f"[~] Resuming download from {initial_size} bytes")

    for attempt in range(max_retries):
        try:
            headers = {}
            if initial_size > 0:
                headers["Range"] = f"bytes={initial_size}-"

            response = requests.get(url, headers=headers, stream=True, timeout=60)
            response.raise_for_status()

            # Get total file size
            total_size = initial_size
            if "content-length" in response.headers:
                total_size = initial_size + int(response.headers["content-length"])
            elif "content-range" in response.headers:
                content_range = response.headers["content-range"]
                if "/" in content_range:
                    total_size = int(content_range.split("/")[-1])

            print(f"[~] Total size: {total_size / (1024**3):.2f} GB")

            start_time = time.time()
            last_progress_time = start_time
            mode = "ab" if initial_size > 0 else "wb"

            with open(dest, mode) as f:
                downloaded = initial_size
                for chunk in response.iter_content(chunk_size=chunk_size):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)

                        current_time = time.time()
                        if current_time - last_progress_time >= 10:
                            if total_size > 0:
                                progress = (downloaded / total_size) * 100
                                speed = downloaded / (current_time - start_time) / (1024**2)
                                print(f"[~] Progress: {progress:.1f}% ({downloaded/(1024**3):.2f}/{total_size/(1024**3):.2f} GB) - Speed: {speed:.1f} MB/s")
                            last_progress_time = current_time

            final_size = os.path.getsize(dest)
            if total_size > 0 and final_size < total_size:
                print(f"[!] Download incomplete: {final_size}/{total_size} bytes")
                initial_size = final_size
                continue

            elapsed_time = time.time() - start_time
            print(f"[+] Download completed in {elapsed_time:.2f} seconds")
            return

        except requests.exceptions.RequestException as e:
            print(f"[!] Download attempt {attempt + 1} failed: {str(e)}")
            if attempt < max_retries - 1:
                print(f"[~] Retrying in 5 seconds...")
                time.sleep(5)
                initial_size = os.path.getsize(dest) if os.path.exists(dest) else 0
            else:
                print(f"[ERROR] Failed to download after {max_retries} attempts")
                raise

def extract_tar_file(tar_path: str, extract_dir: str) -> None:
    """Extract tar file with error handling"""
    print(f"[~] Extracting {tar_path} to {extract_dir}")
    try:
        import tarfile
        with tarfile.open(tar_path, "r") as tar:
            tar.extractall(path=extract_dir)
        print(f"[+] Extraction completed")
        os.remove(tar_path)
        print(f"[+] Cleaned up {tar_path}")
    except Exception as e:
        print(f"[ERROR] Failed to extract {tar_path}: {str(e)}")
        raise

def download_weights(url: str, dest: str) -> None:
    """Download and extract model weights with robust error handling."""
    start_time = time.time()
    try:
        is_tar = dest.endswith(".tar")
        extract_dir = dest.replace(".tar", "") if is_tar else dest

        if is_tar and os.path.exists(extract_dir) and os.path.isdir(extract_dir):
            print(f"[~] Model already exists at {extract_dir}, skipping download")
            return

        download_with_resume(url, dest)

        if is_tar:
            extract_tar_file(dest, os.path.dirname(dest))

        total_time = time.time() - start_time
        print(f"[+] Model setup completed in {total_time:.2f} seconds")
    except Exception as e:
        print(f"[ERROR] Failed to download weights from {url}: {str(e)}")
        if os.path.exists(dest):
            os.remove(dest)
        raise

def decode_base64_file(base64_string: str, output_path: str) -> str:
    """Decode base64 string and save to file"""
    try:
        if "," in base64_string:
            base64_string = base64_string.split(",")[1]
        with open(output_path, "wb") as f:
            f.write(base64.b64decode(base64_string))
        return output_path
    except Exception as e:
        raise ValueError(f"Failed to decode base64 file: {str(e)}")

def encode_file_to_base64(file_path: str) -> str:
    """Encode file to base64 string"""
    try:
        with open(file_path, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")
    except Exception as e:
        raise ValueError(f"Failed to encode file to base64: {str(e)}")

# Download models function
@app.function(
    volumes={MODEL_DIR.as_posix(): volume},
    image=download_image,
    secrets=[huggingface_secret],
    timeout=3600,  # 1 hour timeout for downloads
)
def download_models():
    """Download all required MultiTalk models"""
    MODEL_CACHE = MODEL_DIR / "weights"
    BASE_URL = f"https://weights.replicate.delivery/default/multitalk/{MODEL_CACHE.name}/"
    
    # Create model cache directory
    MODEL_CACHE.mkdir(parents=True, exist_ok=True)
    
    # Download model weights
    model_files = [
        "chinese-wav2vec2-base.tar",
        "MeiGen-MultiTalk.tar", 
        "Wan2.1-I2V-14B-480P.tar",
    ]
    
    print("[~] Starting model downloads...")
    for model_file in model_files:
        url = BASE_URL + model_file
        dest_path = MODEL_CACHE / model_file
        download_weights(url, str(dest_path))
    
    # Commit changes to volume
    volume.commit()
    print("✅ All models downloaded and committed to volume!")

# Inference class
class MultiTalkInference:
    def __init__(self):
        """Initialize the MultiTalk inference handler"""
        import torch
        from transformers import Wav2Vec2FeatureExtractor
        # Note: You'll need to add proper imports for MultiTalk components
        # import wan
        # from wan.configs import WAN_CONFIGS
        # from src.audio_analysis.wav2vec2 import Wav2Vec2Model
        # from wan.utils.multitalk_utils import save_video_ffmpeg
        
        # Model paths
        self.model_cache = MODEL_DIR / "weights"
        self.ckpt_dir = self.model_cache / "Wan2.1-I2V-14B-480P"
        self.wav2vec_dir = self.model_cache / "chinese-wav2vec2-base"
        self.multitalk_dir = self.model_cache / "MeiGen-MultiTalk"
        
        # Initialize device
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        
        # Load wav2vec models
        print("Loading wav2vec models...")
        audio_device = (
            self.device
            if torch.cuda.is_available()
            and torch.cuda.get_device_properties(0).total_memory > 40 * 1024**3
            else "cpu"
        )
        print(f"Loading audio encoder on: {audio_device}")
        
        # Note: Replace with actual imports when available
        # self.audio_encoder = Wav2Vec2Model.from_pretrained(
        #     str(self.wav2vec_dir), local_files_only=True, attn_implementation="eager"
        # ).to(audio_device)
        # self.audio_encoder.feature_extractor._freeze_parameters()
        # self.wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
        #     str(self.wav2vec_dir), local_files_only=True
        # )
        self.audio_device = audio_device
        
        # Load MultiTalk pipeline
        print("Loading MultiTalk pipeline...")
        # self.cfg = WAN_CONFIGS["multitalk-14B"]
        # self.wan_i2v = wan.MultiTalkPipeline(
        #     config=self.cfg,
        #     checkpoint_dir=str(self.ckpt_dir),
        #     device_id=0,
        #     rank=0,
        #     t5_fsdp=False,
        #     dit_fsdp=False,
        #     use_usp=False,
        #     t5_cpu=False,
        # )
        
        # GPU optimizations
        if torch.cuda.is_available():
            vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            print(f"🔍 Detected {vram_gb:.1f}GB VRAM")
            
            if vram_gb > 40:
                print("🚀 High-VRAM detected: Enabling maximum performance optimizations")
                torch.backends.cuda.enable_flash_sdp(True)
                torch.backends.cudnn.benchmark = True
                torch.backends.cuda.matmul.allow_tf32 = True
                torch.backends.cudnn.allow_tf32 = True
                torch.cuda.empty_cache()
                print("⚡ Enabled Flash-SDP, TF32, and cuDNN benchmarking")
            else:
                print("🔧 Standard GPU optimizations enabled")
                torch.backends.cuda.enable_flash_sdp(True)
                torch.cuda.empty_cache()
        
        print("✅ Model setup completed successfully!")
    
    def generate_video(self, request: MultiTalkRequest) -> MultiTalkResponse:
        """Generate video from the request"""
        import torch
        import numpy as np
        import random
        import soundfile as sf
        
        try:
            # Validate parameters
            if not isinstance(request.num_frames, int) or request.num_frames < 25 or request.num_frames > 201:
                raise ValueError("num_frames must be an integer between 25 and 201")
            if not isinstance(request.sampling_steps, int) or request.sampling_steps < 2 or request.sampling_steps > 100:
                raise ValueError("sampling_steps must be an integer between 2 and 100")
            
            # Auto-correct frame count
            num_frames = request.num_frames
            original_frames = num_frames
            if (num_frames - 1) % 4 != 0:
                n_lower = (num_frames - 1) // 4
                n_upper = n_lower + 1
                frames_lower = 4 * n_lower + 1
                frames_upper = 4 * n_upper + 1
                
                if abs(num_frames - frames_lower) <= abs(num_frames - frames_upper):
                    num_frames = frames_lower
                else:
                    num_frames = frames_upper
                
                num_frames = max(25, min(num_frames, 201))
                
                while (num_frames - 1) % 4 != 0 and num_frames <= 201:
                    num_frames += 1
                
                print(f"📐 Auto-corrected num_frames from {original_frames} to {num_frames}")
            
            # Set seed
            if request.seed is None:
                seed = random.randint(0, 99999999)
            else:
                seed = request.seed
            
            print(f"🎬 Generating video with seed: {seed}")
            
            with tempfile.TemporaryDirectory() as temp_dir:
                # Save input files
                image_path = os.path.join(temp_dir, "input_image.jpg")
                first_audio_path = os.path.join(temp_dir, "first_audio.wav")
                
                decode_base64_file(request.image, image_path)
                decode_base64_file(request.first_audio, first_audio_path)
                
                second_audio_path = None
                if request.second_audio:
                    second_audio_path = os.path.join(temp_dir, "second_audio.wav")
                    decode_base64_file(request.second_audio, second_audio_path)
                
                is_multi_person = second_audio_path is not None
                audio_save_dir = os.path.join(temp_dir, "audio_embeddings")
                os.makedirs(audio_save_dir, exist_ok=True)
                
                # Process audio and generate embeddings
                if is_multi_person:
                    print("🎤 Processing multi-person audio...")
                    audio_type = "add"
                    speech1, speech2, combined_speech = audio_prepare_multi(
                        first_audio_path, second_audio_path, audio_type
                    )
                    
                    # Generate embeddings
                    embedding1 = get_embedding(
                        speech1, self.wav2vec_feature_extractor, self.audio_encoder, 
                        device=self.audio_device
                    )
                    embedding2 = get_embedding(
                        speech2, self.wav2vec_feature_extractor, self.audio_encoder,
                        device=self.audio_device
                    )
                    
                    # Save embeddings and audio
                    emb1_path = os.path.join(audio_save_dir, "1.pt")
                    emb2_path = os.path.join(audio_save_dir, "2.pt")
                    sum_audio_path = os.path.join(audio_save_dir, "sum.wav")
                    
                    torch.save(embedding1, emb1_path)
                    torch.save(embedding2, emb2_path)
                    sf.write(sum_audio_path, combined_speech, 16000)
                    
                    input_data = {
                        "prompt": request.prompt,
                        "cond_image": image_path,
                        "audio_type": audio_type,
                        "cond_audio": {"person1": emb1_path, "person2": emb2_path},
                        "video_audio": sum_audio_path,
                    }
                else:
                    print("🎤 Processing single-person audio...")
                    speech = audio_prepare_single(first_audio_path)
                    embedding = get_embedding(
                        speech, self.wav2vec_feature_extractor, self.audio_encoder,
                        device=self.audio_device
                    )
                    
                    emb_path = os.path.join(audio_save_dir, "1.pt")
                    sum_audio_path = os.path.join(audio_save_dir, "sum.wav")
                    
                    torch.save(embedding, emb_path)
                    sf.write(sum_audio_path, speech, 16000)
                    
                    input_data = {
                        "prompt": request.prompt,
                        "cond_image": image_path,
                        "cond_audio": {"person1": emb_path},
                        "video_audio": sum_audio_path,
                    }
                
                print("🎬 Generating video...")
                
                # Configure generation parameters
                high_vram = (
                    torch.cuda.is_available()
                    and torch.cuda.get_device_properties(0).total_memory > 40 * 1024**3
                )
                
                if request.turbo:
                    teacache_thresh = 0.8
                    text_guide_scale = 3.0
                    audio_guide_scale = 3.0
                    shift = 5.0
                    offload_model = False
                else:
                    teacache_thresh = 0.3
                    text_guide_scale = 5.0
                    audio_guide_scale = 4.0
                    shift = 7.0
                    offload_model = not high_vram
                
                extra_args = SimpleNamespace(
                    use_teacache=True,
                    teacache_thresh=teacache_thresh,
                    use_apg=False,
                    size="multitalk-480",
                )
                
                # Generate video (replace with actual generation call)
                # video = self.wan_i2v.generate(
                #     input_data,
                #     size_buckget="multitalk-480",
                #     motion_frame=25,
                #     frame_num=num_frames,
                #     shift=shift,
                #     sampling_steps=request.sampling_steps,
                #     text_guide_scale=text_guide_scale,
                #     audio_guide_scale=audio_guide_scale,
                #     seed=seed,
                #     offload_model=offload_model,
                #     max_frames_num=num_frames,
                #     extra_args=extra_args,
                # )
                
                # Save video (replace with actual save call)
                output_name = f"multitalk_{abs(hash(request.prompt + str(seed))) % 10000}"
                # save_video_ffmpeg(video, output_name, [input_data["video_audio"]])
                
                # For now, create a dummy video file
                output_file = f"{output_name}.mp4"
                with open(output_file, "wb") as f:
                    f.write(b"dummy video content")
                
                # Encode video to base64
                video_b64 = encode_file_to_base64(output_file)
                
                # Cleanup
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                
                print("✅ Video generation completed successfully!")
                
                return MultiTalkResponse(
                    video=video_b64,
                    metadata={
                        "seed": seed,
                        "num_frames": num_frames,
                        "sampling_steps": request.sampling_steps,
                        "turbo": request.turbo,
                        "is_multi_person": is_multi_person,
                        "prompt": request.prompt,
                        "original_frames_requested": original_frames,
                    }
                )
        
        except Exception as e:
            print(f"❌ Error during inference: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

# Inference endpoint
@app.function(
    volumes={MODEL_DIR.as_posix(): volume},
    image=base_image,
    gpu=modal.gpu.A100(size="40GB"),  # Adjust GPU as needed
    secrets=[huggingface_secret],
    timeout=1800,  # 30 minute timeout
    container_idle_timeout=300,  # 5 minute idle timeout
)
@modal.web_endpoint(method="POST")
def generate_video(request: MultiTalkRequest) -> Union[MultiTalkResponse, ErrorResponse]:
    """Main inference endpoint"""
    try:
        # Initialize inference handler
        handler = MultiTalkInference()
        
        # Generate video
        response = handler.generate_video(request)
        
        return response
        
    except Exception as e:
        print(f"❌ Endpoint error: {str(e)}")
        return ErrorResponse(error=str(e), type=type(e).__name__)