import io
import os
import asyncio
from typing import Optional

import torch
import torchaudio as ta
import torchaudio.functional as F
from chatterbox.tts import ChatterboxTTS
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from concurrent.futures import ThreadPoolExecutor


class TTSConfig:
    """Configuration for TTS model and processing"""
    MODEL_REPO = "grandhigh/Chatterbox-TTS-Indonesian"
    CHECKPOINT = "t3_cfg.safetensors"
    DEVICE = "cpu"
    
    # Generation parameters
    TEMPERATURE = 0.65
    TOP_P = 0.88
    REPETITION_PENALTY = 1.25
    
    # Audio processing
    AUDIO_GAIN_DB = 0.8


class AudioProcessor:
    """Audio enhancement utilities"""
    
    @staticmethod
    def generate_pink_noise(shape, device):
        """Generate pink noise for audio enhancement"""
        white = torch.randn(shape, device=device)
        pink = torch.zeros_like(white)
        b = torch.zeros(7)

        if len(shape) == 1:
            for j in range(shape[0]):
                w = white[j].item()
                b[0] = 0.99886 * b[0] + w * 0.0555179
                b[1] = 0.99332 * b[1] + w * 0.0750759
                b[2] = 0.96900 * b[2] + w * 0.1538520
                b[3] = 0.86650 * b[3] + w * 0.3104856
                b[4] = 0.55000 * b[4] + w * 0.5329522
                b[5] = -0.7616 * b[5] - w * 0.0168980
                pink[j] = (b[0]+b[1]+b[2]+b[3]+b[4]+b[5]+b[6] + w*0.5362) * 0.11
                b[6] = w * 0.115926
        else:
            for i in range(shape[0]):
                b = torch.zeros(7)
                for j in range(shape[1]):
                    w = white[i, j].item()
                    b[0] = 0.99886 * b[0] + w * 0.0555179
                    b[1] = 0.99332 * b[1] + w * 0.0750759
                    b[2] = 0.96900 * b[2] + w * 0.1538520
                    b[3] = 0.86650 * b[3] + w * 0.3104856
                    b[4] = 0.55000 * b[4] + w * 0.5329522
                    b[5] = -0.7616 * b[5] - w * 0.0168980
                    pink[i, j] = (b[0]+b[1]+b[2]+b[3]+b[4]+b[5]+b[6] + w*0.5362) * 0.11
                    b[6] = w * 0.115926

        return pink * 0.1

    @staticmethod
    def enhance_audio(wav, sr):
        """Apply audio enhancements: normalization, filtering, compression"""
        # Normalize
        peak = wav.abs().max()
        if peak > 0:
            wav = wav / (peak + 1e-8) * 0.95

        # Apply filters
        wav = F.highpass_biquad(wav, sr, cutoff_freq=60)
        wav = F.lowpass_biquad(wav, sr, cutoff_freq=10000)
        wav = F.bass_biquad(wav, sr, gain=1.5, central_freq=200, Q=0.7)
        wav = F.treble_biquad(wav, sr, gain=-1.2, central_freq=6000, Q=0.7)

        # Compression
        threshold = 0.6
        ratio = 2.5
        abs_wav = wav.abs()
        compressed = wav.clone()
        mask = abs_wav > threshold
        compressed[mask] = torch.sign(wav[mask]) * (threshold + (abs_wav[mask] - threshold) / ratio)

        wav = compressed
        wav = torch.tanh(wav * 1.08)
        
        # Add pink noise
        wav = wav + AudioProcessor.generate_pink_noise(wav.shape, wav.device) * 0.0003
        wav = F.gain(wav, gain_db=TTSConfig.AUDIO_GAIN_DB)

        # Final normalization
        peak = wav.abs().max()
        if peak > 0:
            wav = wav / peak * 0.88

        return wav

    @staticmethod
    def save_tensor_to_wav(wav_tensor: torch.Tensor, sr: int, out_wav_path: str):
        """Save a torch tensor to WAV file"""
        # Ensure float32 CPU tensor
        if wav_tensor.device.type != "cpu":
            wav_tensor = wav_tensor.cpu()
        if wav_tensor.dtype != torch.float32:
            wav_tensor = wav_tensor.type(torch.float32)

        # torchaudio.save requires shape [channels, samples]
        if wav_tensor.dim() == 1:
            wav_out = wav_tensor.unsqueeze(0)
        else:
            wav_out = wav_tensor

        # Save directly as WAV
        ta.save(out_wav_path, wav_out, sr, format="wav")

    @staticmethod
    def tensor_to_wav_buffer(wav_tensor: torch.Tensor, sr: int) -> io.BytesIO:
        """Convert torch tensor to WAV buffer"""
        buf = io.BytesIO()
        if wav_tensor.dim() == 1:
            wav_out = wav_tensor.unsqueeze(0)
        else:
            wav_out = wav_tensor
        ta.save(buf, wav_out, sr, format="wav")
        buf.seek(0)
        return buf


class TTSEngine:
    """Main TTS engine with model management"""
    
    def __init__(self, config: TTSConfig, thread_pool: ThreadPoolExecutor):
        self.config = config
        self.thread_pool = thread_pool
        self.model = None
        self.model_lock = asyncio.Lock()
        self.sr = None
        
    def load_model(self):
        """Load the TTS model and checkpoint"""
        print("Loading model...")
        self.model = ChatterboxTTS.from_pretrained(device=self.config.DEVICE)
        ckpt = hf_hub_download(repo_id=self.config.MODEL_REPO, filename=self.config.CHECKPOINT)
        state = load_file(ckpt, device=self.config.DEVICE)

        self.model.t3.to(self.config.DEVICE).load_state_dict(state)
        self.model.t3.eval()
        
        # Disable dropout
        for m in self.model.t3.modules():
            if hasattr(m, "training"):
                m.training = False
            if isinstance(m, torch.nn.Dropout):
                m.p = 0
        
        self.sr = self.model.sr
        print("Model ready.")
    
    async def generate(self, text: str, audio_prompt_path: str) -> torch.Tensor:
        """Generate audio from text with voice prompt"""
        async with self.model_lock:
            def blocking_generate():
                with torch.no_grad():
                    return self.model.generate(
                        text,
                        audio_prompt_path=audio_prompt_path,
                        temperature=self.config.TEMPERATURE,
                        top_p=self.config.TOP_P,
                        repetition_penalty=self.config.REPETITION_PENALTY,
                    )
            
            wav = await asyncio.get_event_loop().run_in_executor(
                self.thread_pool, 
                blocking_generate
            )
            return wav
    
    async def generate_and_enhance(self, text: str, audio_prompt_path: str) -> torch.Tensor:
        """Generate and enhance audio"""
        wav = await self.generate(text, audio_prompt_path)
        
        # Enhance audio (CPU-bound)
        wav = await asyncio.get_event_loop().run_in_executor(
            self.thread_pool,
            lambda: AudioProcessor.enhance_audio(wav.cpu(), self.sr)
        )
        
        return wav
    
    async def generate_to_file(self, text: str, audio_prompt_path: str, output_path: str):
        """Generate audio and save to file"""
        wav = await self.generate_and_enhance(text, audio_prompt_path)
        
        # Save to WAV
        await asyncio.get_event_loop().run_in_executor(
            self.thread_pool,
            AudioProcessor.save_tensor_to_wav,
            wav,
            self.sr,
            output_path
        )
    
    async def generate_to_buffer(self, text: str, audio_prompt_path: str) -> io.BytesIO:
        """Generate audio and return as WAV buffer"""
        wav = await self.generate_and_enhance(text, audio_prompt_path)
        
        # Convert to buffer
        return AudioProcessor.tensor_to_wav_buffer(wav, self.sr)