import os
import torch
import logging
from typing import Optional, Dict, Any
from fastapi import FastAPI, HTTPException, status, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from starlette.background import BackgroundTask
from pydantic import BaseModel
import torchaudio
import io
import tempfile
import numpy as np
from scipy.io import wavfile
import requests
import soundfile as sf
import subprocess
import imageio_ffmpeg

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="YarnGPT2 TTS API",
    description="Text-to-Speech API using YarnGPT2 model for Nigerian accents and languages",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

ASK_URL = os.getenv("ASK_URL", "https://remostart-milestone-one-farmlingua-ai.hf.space/ask")

asr_models = {
    "ha": {"repo": "NCAIR1/Hausa-ASR", "model": None, "proc": None},
    "yo": {"repo": "NCAIR1/Yoruba-ASR", "model": None, "proc": None},
    "ig": {"repo": "NCAIR1/Igbo-ASR", "model": None, "proc": None},
    "en": {"repo": "NCAIR1/NigerianAccentedEnglish", "model": None, "proc": None},
}

model = None
audio_tokenizer = None
device = None

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")

class TTSRequest(BaseModel):
    text: str
    language: str = "english"
    speaker_name: str = "idera"
    temperature: float = 0.1
    repetition_penalty: float = 1.1
    max_length: int = 4000

class TTSResponse(BaseModel):
    message: str
    audio_url: str


class SpeakRequest(BaseModel):
    text: str
    language: str
    speaker_name: str | None = None
    temperature: float | None = 0.1
    repetition_penalty: float | None = 1.1
    max_length: int | None = 4000

def load_audio_tokenizer():
    """
    Load the AudioTokenizerV2 for processing.
    """
    global audio_tokenizer
    
    try:
        config_paths = [
            "./wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml",
            "./models/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml",
            "./wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
        ]
        
        model_paths = [
            "./wavtokenizer_large_speech_320_24k.ckpt",
            "./models/wavtokenizer_large_speech_320_24k.ckpt",
            "./wavtokenizer_large_speech_320_24k.ckpt"
        ]
        
        config_path = next((p for p in config_paths if os.path.exists(p)), config_paths[0])
        
        model_path = None
        for mp in model_paths:
            if os.path.exists(mp):
                model_path = mp
                break
        
        if not model_path or not os.path.exists(model_path):
            logger.warning("Checkpoint file not found, attempting to download...")
            try:
                import subprocess, tempfile, shutil as _shutil
                target_dir = os.environ.get("MODEL_DIR", "./models")
                os.makedirs(target_dir, exist_ok=True)
                tmp_file = os.path.join("/tmp", f"wavtokenizer_large_speech_320_24k.ckpt.{os.getpid()}.part")
                result = subprocess.run([
                    "gdown", "--fuzzy", "1-ASeEkrn4HY49yZWHTASgfGFNXdVnLTt", 
                    "-O", tmp_file
                ], check=False, capture_output=True, text=True, env=os.environ.copy())
                final_path = os.path.join(target_dir, "wavtokenizer_large_speech_320_24k.ckpt")
                if result.returncode == 0 and os.path.exists(tmp_file):
                    _shutil.move(tmp_file, final_path)
                    model_path = final_path
                    logger.info("Checkpoint downloaded successfully")
                else:
                    model_path = model_paths[0]
                    logger.warning(f"Checkpoint download failed: {result.stderr}, using fallback path")
            except Exception as e:
                logger.warning(f"Could not download checkpoint: {e}, using fallback path")
                model_path = model_paths[0]
        
        from yarngpt.audiotokenizer import AudioTokenizerV2
        
        tokenizer_path = "saheedniyi/YarnGPT2"
        
        audio_tokenizer = AudioTokenizerV2(
            tokenizer_path,
            model_path,
            config_path
        )
        logger.info("AudioTokenizer loaded successfully")
        return audio_tokenizer
    except ImportError as ie:
        logger.warning(f"yarngpt package not found: {ie}")
        try:
            from transformers import AutoTokenizer
            
            tokenizer_path = "saheedniyi/YarnGPT2"
            
            class AudioTokenizerWrapper:
                def __init__(self, tokenizer_path):
                    self.tokenizer_path = tokenizer_path
                    self.device = device
                    self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
                    logger.info("Using fallback tokenizer")
                    
                def create_prompt(self, text, lang="english", speaker_name="idera"):
                    """
                    Create a prompt.
                    """
                    speaker_tag = f"<{speaker_name}>"
                    lang_tag = f"<{lang}>"
                    return f"{speaker_tag}{lang_tag}{text}</s>"
                
                def tokenize_prompt(self, prompt):
                    """
                    Tokenize prompt.
                    """
                    return self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
                
                def get_codes(self, output):
                    """
                    Extract audio codes.
                    """
                    return output
                
                def get_audio(self, codes):
                    """
                    Convert codes to audio waveform.
                    """
                    # Placeholder implementation
                    import numpy as np
                    sample_rate = 24000
                    duration = 3.0  # Default duration
                    audio = np.random.randn(int(duration * sample_rate)).astype(np.float32)
                    return torch.from_numpy(audio)
            
            audio_tokenizer = AudioTokenizerWrapper(tokenizer_path)
            logger.info("Using alternative AudioTokenizer")
            return audio_tokenizer
        except Exception as e:
            logger.error(f"Failed to load audio tokenizer: {e}")
            raise

def load_model():
    global model
    
    try:
        from transformers import AutoModelForCausalLM
        
        tokenizer_path = "saheedniyi/YarnGPT2"
        logger.info("Loading YarnGPT2 model from HuggingFace...")
        
        model = AutoModelForCausalLM.from_pretrained(
            tokenizer_path,
            torch_dtype="auto"
        ).to(device)
        
        logger.info("YarnGPT2 model loaded successfully")
        return model
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise


def _get_asr(lang_code: str):
    try:
        from transformers import WhisperProcessor, WhisperForConditionalGeneration
    except Exception as e:
        logger.error(f"Transformers missing whisper classes: {e}")
        return None, None

    entry = asr_models.get(lang_code)
    if not entry:
        return None, None
    if entry["model"] is not None and entry["proc"] is not None:
        return entry["model"], entry["proc"]
    repo_id = entry["repo"]
    hf_token = os.getenv("HF_TOKEN")
    try:
        device_t = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Lazy-loading ASR for {lang_code} from {repo_id}...")
        proc = WhisperProcessor.from_pretrained(repo_id, token=hf_token)
        model_asr = WhisperForConditionalGeneration.from_pretrained(repo_id, token=hf_token)
        model_asr.to(device_t)
        model_asr.eval()
        entry["model"], entry["proc"] = model_asr, proc
        return model_asr, proc
    except Exception as e:
        logger.error(f"Failed to load ASR for {lang_code}: {e}")
        entry["model"], entry["proc"] = None, None
        return None, None


def _preprocess_audio_ffmpeg(audio_bytes: bytes, target_sr: int = 16000) -> np.ndarray:
    try:
        with tempfile.NamedTemporaryFile(suffix='.input', delete=False) as in_file:
            in_file.write(audio_bytes)
            in_path = in_file.name
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as out_file:
            out_path = out_file.name

        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
        subprocess.run([
            ffmpeg_exe, '-y', '-i', in_path,
            '-ac', '1',
            '-ar', str(target_sr),
            out_path
        ], check=True, capture_output=True)

        with open(out_path, 'rb') as f:
            wav_data = f.read()

        os.unlink(in_path)
        os.unlink(out_path)

        audio_array, sr = sf.read(io.BytesIO(wav_data))
        if audio_array.ndim > 1:
            audio_array = np.mean(audio_array, axis=1)
        if sr != target_sr:
            ratio = target_sr / sr
            new_len = int(len(audio_array) * ratio)
            audio_array = np.interp(
                np.linspace(0, len(audio_array), new_len),
                np.arange(len(audio_array)),
                audio_array
            )
        audio_array = np.clip(audio_array, -0.99, 0.99)
        audio_array = audio_array - float(np.mean(audio_array))
        return audio_array.astype(np.float32)
    except Exception as e:
        logger.error(f"FFmpeg preprocessing failed: {e}")
        raise HTTPException(status_code=400, detail="Audio preprocessing failed")

@app.on_event("startup")
async def startup_event():
    logger.info("Server started. Models will be loaded on first request.")
    # Don't load models at startup to save storage and startup time
    # Models will be loaded lazily when first API call is made

@app.get("/")
async def root():
    return {
        "name": "YarnGPT2 TTS API",
        "description": "Text-to-Speech API for Nigerian accents and languages",
        "status": "running" if model is not None else "model_loading_failed",
        "available_languages": ["english", "yoruba", "igbo", "hausa"],
        "available_speakers": {
            "english": ["idera", "chinenye", "jude", "emma", "umar", "joke", "zainab", "osagie", "remi", "tayo"],
            "yoruba": ["yoruba_male2", "yoruba_female2", "yoruba_female1"],
            "igbo": ["igbo_female2", "igbo_male2", "igbo_female1"],
            "hausa": ["hausa_female1", "hausa_female2", "hausa_male2", "hausa_male1"]
        }
    }

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy" if model is not None else "degraded",
        "device": str(device),
        "model_loaded": model is not None,
        "tokenizer_loaded": audio_tokenizer is not None
    }


@app.post("/ask")
async def ask(query: str = Form(...)):
    try:
        resp = requests.post(ASK_URL, json={"query": query}, timeout=30)
        resp.raise_for_status()
        return resp.json()
    except Exception as e:
        logger.error(f"ASK error: {e}")
        raise HTTPException(status_code=502, detail="Ask backend error")


@app.post("/transcribe")
async def transcribe(audio_file: UploadFile = File(...), language: str = Form(...)):
    if not audio_file.content_type or not audio_file.content_type.startswith('audio/'):
        raise HTTPException(status_code=400, detail="File must be an audio file")
    if language not in ["yo", "ha", "ig", "en"]:
        raise HTTPException(status_code=400, detail="Language must be one of: yo, ha, ig, en")
    audio_bytes = await audio_file.read()
    audio_array = _preprocess_audio_ffmpeg(audio_bytes)
    model_asr, proc = _get_asr(language)
    if model_asr is None or proc is None:
        raise HTTPException(status_code=500, detail="ASR model not available")
    try:
        device_t = next(model_asr.parameters()).device
        inputs = proc(audio_array, sampling_rate=16000, return_tensors="pt")
        input_features = inputs.input_features.to(device_t)
        with torch.no_grad():
            pred_ids = model_asr.generate(input_features)
        text_list = proc.batch_decode(pred_ids, skip_special_tokens=True)
        transcript = text_list[0].strip() if text_list else ""
        return {"language": language, "transcription": transcript}
    except Exception as e:
        logger.error(f"ASR inference failed: {e}")
        raise HTTPException(status_code=500, detail="ASR inference failed")


@app.post("/speak-ai")
async def speak_ai(audio_file: UploadFile = File(...), language: str = Form(...)):
    # Transcribe
    tr = await transcribe(audio_file=audio_file, language=language)
    query = tr.get("transcription", "")
    # Ask
    try:
        ans = requests.post(ASK_URL, json={"query": query}, timeout=30)
        ans.raise_for_status()
        answer_text = ans.json().get("answer", query) or query
    except Exception as e:
        logger.warning(f"Ask failed ({e}); falling back to transcript")
        answer_text = query
    # TTS via YarnGPT2 /speak logic
    speak_req = SpeakRequest(text=answer_text, language=_map_lang_code(language))
    return await speak(speak_req)


def _map_lang_code(code: str) -> str:
    m = {"yo": "yoruba", "ha": "hausa", "ig": "igbo", "en": "english"}
    return m.get(code.lower(), "english")

@app.post("/tts")
async def text_to_speech(request: TTSRequest):
    # Lazy load models on first request
    global model, audio_tokenizer
    if model is None:
        logger.info("Loading YarnGPT2 model (lazy loading)...")
        load_model()
    if audio_tokenizer is None:
        logger.info("Loading audio tokenizer (lazy loading)...")
        load_audio_tokenizer()
    
    if model is None or audio_tokenizer is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Model loading failed. Please check logs."
        )


@app.post("/speak")
async def speak(request: SpeakRequest):
    global model, audio_tokenizer
    if model is None:
        logger.info("Loading YarnGPT2 model (lazy loading)...")
        load_model()
    if audio_tokenizer is None:
        logger.info("Loading audio tokenizer (lazy loading)...")
        load_audio_tokenizer()

    if model is None or audio_tokenizer is None:
        raise HTTPException(status_code=503, detail="Model loading failed. Please check logs.")

    default_speakers = {
        "english": "idera",
        "yoruba": "yoruba_male2",
        "igbo": "igbo_male2",
        "hausa": "hausa_female1",
    }
    language = request.language.lower().strip()
    speaker = request.speaker_name or default_speakers.get(language, "idera")

    try:
        prompt = audio_tokenizer.create_prompt(
            request.text,
            lang=language,
            speaker_name=speaker,
        )
        input_ids = audio_tokenizer.tokenize_prompt(prompt)
        with torch.no_grad():
            output = model.generate(
                input_ids=input_ids,
                temperature=request.temperature or 0.1,
                repetition_penalty=request.repetition_penalty or 1.1,
                max_length=request.max_length or 4000,
            )

        codes = audio_tokenizer.get_codes(output)
        audio = audio_tokenizer.get_audio(codes)

        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
        if isinstance(audio, torch.Tensor):
            audio_np = audio.detach().cpu().numpy()
        else:
            audio_np = np.asarray(audio)
        if audio_np.ndim == 2:
            audio_np = audio_np[0]
        audio_np = np.clip(audio_np, -1.0, 1.0)
        audio_int16 = (audio_np * 32767.0).astype(np.int16)
        wavfile.write(temp_file.name, 24000, audio_int16)

        return FileResponse(
            temp_file.name,
            media_type="audio/wav",
            filename="speech.wav",
            background=BackgroundTask(lambda: os.path.exists(temp_file.name) and os.unlink(temp_file.name))
        )
    except Exception as e:
        logger.error(f"Speak error: {e}")
        raise HTTPException(status_code=500, detail=f"Speak failed: {e}")

@app.post("/tts-stream")
async def text_to_speech_stream(request: TTSRequest):
    global model, audio_tokenizer
    if model is None:
        logger.info("Loading YarnGPT2 model (lazy loading)...")
        load_model()
    if audio_tokenizer is None:
        logger.info("Loading audio tokenizer (lazy loading)...")
        load_audio_tokenizer()
    
    if model is None or audio_tokenizer is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Model loading failed. Please check logs."
        )
    
    try:
        # Create prompt
        prompt = audio_tokenizer.create_prompt(
            request.text,
            lang=request.language,
            speaker_name=request.speaker_name
        )
        
        # Tokenize
        input_ids = audio_tokenizer.tokenize_prompt(prompt)
        
        # Generate
        logger.info(f"Generating speech (streaming) for text: {request.text[:50]}...")
        with torch.no_grad():
            output = model.generate(
                input_ids=input_ids,
                temperature=request.temperature,
                repetition_penalty=request.repetition_penalty,
                max_length=request.max_length,
            )
        
        # Get audio
        codes = audio_tokenizer.get_codes(output)
        audio = audio_tokenizer.get_audio(codes)
        
        # Convert to bytes using scipy
        buffer = io.BytesIO()
        if isinstance(audio, torch.Tensor):
            audio_np = audio.detach().cpu().numpy()
        else:
            audio_np = np.asarray(audio)
        if audio_np.ndim == 2:
            audio_np = audio_np[0]
        audio_np = np.clip(audio_np, -1.0, 1.0)
        audio_int16 = (audio_np * 32767.0).astype(np.int16)
        wavfile.write(buffer, 24000, audio_int16)
        buffer.seek(0)
        
        return StreamingResponse(
            buffer,
            media_type="audio/wav",
            headers={"Content-Disposition": "attachment; filename=speech.wav"},
            background=BackgroundTask(buffer.close)
        )
        
    except Exception as e:
        logger.error(f"Error generating speech: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to generate speech: {str(e)}"
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)