import os
import io
import tempfile
import subprocess
import requests
import torch
import numpy as np
import soundfile as sf
from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM
import imageio_ffmpeg
import logging
from contextlib import asynccontextmanager
import uvicorn
import nest_asyncio

nest_asyncio.apply()

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

@asynccontextmanager
async def lifespan(app: FastAPI):
    load_models()
    yield

app = FastAPI(title="Farmlingua AI Speech Interface", version="1.0.0", lifespan=lifespan)

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


ASK_URL = "https://remostart-milestone-one-farmlingua-ai.hf.space/ask"
tts_ha, tts_en, tts_yo, tts_ig = None, None, None, None
natlas_tokenizer, natlas_model = None, None

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},
}

def load_models():
    global tts_ha, tts_en, tts_yo, tts_ig, natlas_tokenizer, natlas_model
    device = 0 if torch.cuda.is_available() else -1
    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        hf_token = hf_token.strip()
    if not hf_token:
        logger.warning("HF_TOKEN not set! This may cause authentication failures for gated repositories.")
        logger.warning("Please set HF_TOKEN environment variable to access restricted models.")
    else:
        logger.info("HF_TOKEN is set and ready for authenticated model access.")
    
    logger.info("Loading N-ATLaS language identification model...")
    _load_natlas()
    
    logger.info("Loading TTS models...")
    try:
        tts_ha = pipeline("text-to-speech", model="facebook/mms-tts-hau", device=device)
        logger.info("Loaded TTS (Hausa)")
    except Exception as e:
        logger.exception("Failed to load TTS (Hausa)")
        tts_ha = None
    try:
        tts_en = pipeline("text-to-speech", model="facebook/mms-tts-eng", device=device)
        logger.info("Loaded TTS (English)")
    except Exception:
        logger.exception("Failed to load TTS (English)")
        tts_en = None
    try:
        tts_yo = pipeline("text-to-speech", model="facebook/mms-tts-yor", device=device)
        logger.info("Loaded TTS (Yoruba)")
    except Exception:
        logger.exception("Failed to load TTS (Yoruba)")
        tts_yo = None
   
    tts_ig = None
    logger.info("Igbo TTS model disabled - will return text responses for Igbo language")

    logger.info("Deferred ASR model loads: will lazy-load per language on first use")

def _get_asr(lang_code: str):
    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")
    if hf_token:
        hf_token = hf_token.strip()
    try:
        device = 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 = WhisperForConditionalGeneration.from_pretrained(repo_id, token=hf_token)
        model.to(device)
        model.eval()
        entry["model"], entry["proc"] = model, proc
        logger.info(f"Loaded ASR for {lang_code}")
        return model, proc
    except Exception:
        logger.exception(f"Failed to load ASR for {lang_code} ({repo_id})")
        entry["model"], entry["proc"] = None, None
        return None, None

def _run_whisper(model: WhisperForConditionalGeneration, proc: WhisperProcessor, audio_array: np.ndarray) -> str:
    try:
        device = next(model.parameters()).device
        inputs = proc(audio_array, sampling_rate=16000, return_tensors="pt")
        input_features = inputs.input_features.to(device)
        with torch.no_grad():
            predicted_ids = model.generate(input_features)
        text_list = proc.batch_decode(predicted_ids, skip_special_tokens=True)
        return text_list[0] if text_list else ""
    except Exception:
        logging.exception("Whisper ASR inference failed")
        return ""

def preprocess_audio_ffmpeg(audio_data: bytes, target_sr: int = 16000) -> np.ndarray:
    try:
        with tempfile.NamedTemporaryFile(suffix='.input', delete=False) as in_file:
            in_file.write(audio_data)
            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 len(audio_array.shape) > 1:
            audio_array = np.mean(audio_array, axis=1)
        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. Ensure ffmpeg is installed.")

def speech_to_text(audio_data: bytes) -> str:
    audio_array = preprocess_audio_ffmpeg(audio_data)
    candidates = []
    for code in ["yo", "ha", "ig", "en"]:
        model, proc = _get_asr(code)
        if model is None or proc is None:
            continue
        text = _run_whisper(model, proc, audio_array)
        if text:
            candidates.append((code, text))

    for lang_code, text in candidates:
        det = detect_language(text)
        if lang_code == det:
            return text

    if candidates:
        return max((t for _, t in candidates), key=lambda s: len(s or ""))
    return ""


def get_ai_response(text: str) -> str:
    try:
        response = requests.post(ASK_URL, json={"query": text}, timeout=30)
        response.raise_for_status()
        result = response.json()
        return result.get("answer", "Sorry, no answer returned.")
    except Exception as e:
        logger.error(f"AI request error: {e}")
        return f"I'm sorry, I couldn't connect to the AI service. You said: '{text}'."

HAUSA_WORDS = [
    "aikin","manoma","gona","amfanin","yanayi","tsaba","fasaha","bisa","noman","shuka",
    "daji","rani","damina","amfani","bidi'a","noma","bashi","manure","tsiro","gishiri"
]

YORUBA_WORDS = [
    "ilé","ọmọ","òun","awọn","agbẹ","oko","ọgbà","irugbin","àkọsílẹ","omi","ojo","àgbàlá","irọlẹ"
]

IGBO_WORDS = [
    "ugbo","akụkọ","mmiri","ala","ọrụ","ncheta","ọhụrụ","ugwu","nri","ahụhụ"
]

def _load_natlas():
    global natlas_tokenizer, natlas_model
    if natlas_tokenizer is not None and natlas_model is not None:
        logger.info("N-ATLaS model already loaded")
        return True
    
    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        hf_token = hf_token.strip()
    
    if not hf_token:
        logger.error("HF_TOKEN not available for N-ATLaS model access")
        return False
    
    try:
        logger.info("Loading N-ATLaS language identification model...")
        logger.info("Model files are pre-cached from Docker build, loading should be faster...")
        
        natlas_tokenizer = AutoTokenizer.from_pretrained("NCAIR1/N-ATLaS", token=hf_token)
        natlas_model = AutoModelForCausalLM.from_pretrained(
            "NCAIR1/N-ATLaS",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
            token=hf_token,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
            use_cache=True
        )
        logger.info("Successfully loaded N-ATLaS language identification model")
        return True
    except Exception as e:
        logger.exception(f"Failed to load N-ATLaS model: {e}")
        natlas_tokenizer, natlas_model = None, None
        return False

def detect_language(text: str) -> str:
    logger.info(f"Detecting language for text: '{text[:50]}...'")
    
    if not _load_natlas():
        logger.warning("N-ATLaS model not available, falling back to keyword detection")
        text_lower = text.lower()
        if any(word in text_lower for word in HAUSA_WORDS):
            logger.info("Keyword detection: Hausa")
            return "ha"
        elif any(word in text_lower for word in YORUBA_WORDS):
            logger.info("Keyword detection: Yoruba")
            return "yo"
        elif any(word in text_lower for word in IGBO_WORDS):
            logger.info("Keyword detection: Igbo")
            return "ig"
        else:
            logger.info("Keyword detection: English (default)")
            return "en"
    
    try:
        logger.info("Using N-ATLaS for language detection")
        messages = [
            {'role': 'system', 'content': 'You are a language identification assistant. Identify the language of the given text and respond with only the language code: "en" for English, "ha" for Hausa, "yo" for Yoruba, or "ig" for Igbo.'},
            {'role': 'user', 'content': f'What language is this text written in? "{text}"'}
        ]
        
        formatted_text = natlas_tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=False
        )
        
        input_tokens = natlas_tokenizer(formatted_text, return_tensors='pt', add_special_tokens=False)
        if torch.cuda.is_available():
            input_tokens = input_tokens.to('cuda')
        
        with torch.no_grad():
            outputs = natlas_model.generate(
                **input_tokens,
                max_new_tokens=10,
                use_cache=True,
                repetition_penalty=1.1,
                temperature=0.1,
                do_sample=False
            )
        
        response = natlas_tokenizer.batch_decode(outputs)[0]
        response_text = response.split(messages[1]['content'])[-1].strip().lower()
        
        logger.info(f"N-ATLaS response: '{response_text}'")
        
        if 'ha' in response_text:
            logger.info("N-ATLaS detection: Hausa")
            return "ha"
        elif 'yo' in response_text:
            logger.info("N-ATLaS detection: Yoruba")
            return "yo"
        elif 'ig' in response_text:
            logger.info("N-ATLaS detection: Igbo")
            return "ig"
        else:
            logger.info("N-ATLaS detection: English (default)")
            return "en"
            
    except Exception as e:
        logger.exception(f"Language detection failed: {e}")
        logger.warning("Falling back to keyword detection due to N-ATLaS error")
        text_lower = text.lower()
        if any(word in text_lower for word in HAUSA_WORDS):
            return "ha"
        elif any(word in text_lower for word in YORUBA_WORDS):
            return "yo"
        elif any(word in text_lower for word in IGBO_WORDS):
            return "ig"
        else:
            return "en"

def text_to_speech_file(text: str) -> str:
    lang = detect_language(text)
    print(f"Detected language: {lang}")
    
    if lang == "ig":
        logger.info("Igbo language detected - returning text response instead of audio")
        fd, path = tempfile.mkstemp(suffix=".txt")
        os.close(fd)
        with open(path, 'w', encoding='utf-8') as f:
            f.write(text)
        return path
    if lang == "ha":
        tts_model = tts_ha
    elif lang == "yo":
        tts_model = tts_yo
    else:
        tts_model = tts_en
    
    if tts_model is None:
        logger.error(f"TTS model for {lang} is not available")
        raise HTTPException(status_code=500, detail=f"TTS model for {lang} is not available")
    
    speech_output = tts_model(text)
    audio_raw = speech_output["audio"]
    sampling_rate = int(speech_output["sampling_rate"]) 

 
    if isinstance(audio_raw, torch.Tensor):
        audio_np = audio_raw.detach().cpu().numpy()
    else:
        audio_np = np.asarray(audio_raw)

    if audio_np.ndim > 1:
        audio_np = audio_np.reshape(-1)
    audio_np = audio_np.astype(np.float32, copy=False)


    audio_clipped = np.clip(audio_np, -1.0, 1.0)
    audio_int16 = (audio_clipped * 32767.0).astype(np.int16)


    fd, path = tempfile.mkstemp(suffix=".wav")
    os.close(fd)

  
    sf.write(path, audio_int16, sampling_rate, format='WAV', subtype='PCM_16')
    return path


@app.get("/")
async def root():
    return {"status": "ok", "message": "System ready"}

@app.get("/health")
async def health():
    natlas_status = "loaded" if natlas_tokenizer is not None and natlas_model is not None else "not_loaded"
    return {
        "message": "Farmlingua AI Speech Interface is running!",
        "natlas_status": natlas_status,
        "tts_models": {
            "hausa": tts_ha is not None,
            "english": tts_en is not None,
            "yoruba": tts_yo is not None,
            "igbo": False
        }
    }

@app.get("/status")
async def status():
    return {
        "natlas_loaded": natlas_tokenizer is not None and natlas_model is not None,
        "loading_message": "N-ATLaS model is loading shards, please wait..." if natlas_tokenizer is None else "N-ATLaS model is ready"
    }

@app.post("/chat")
async def chat(text: str = Form(...), speak: bool = False, raw: bool = False):
    if not text.strip():
        raise HTTPException(status_code=400, detail="Text cannot be empty")
    final_text = text if raw else get_ai_response(text)
    if speak:
        output_path = text_to_speech_file(final_text)
        lang = detect_language(final_text)
        if lang == "ig":
            return FileResponse(output_path, media_type="text/plain", filename="response.txt")
        else:
            return FileResponse(output_path, media_type="audio/wav", filename="response.wav")
    return {"question": text, "answer": final_text}

@app.post("/speak")
async def speak_to_ai(audio_file: UploadFile = File(...), speak: bool = True):
    if not audio_file.content_type.startswith('audio/'):
        raise HTTPException(status_code=400, detail="File must be an audio file")
    audio_data = await audio_file.read()
    transcription = speech_to_text(audio_data)
    ai_response = get_ai_response(transcription)
    if speak:
        output_path = text_to_speech_file(ai_response)
        lang = detect_language(ai_response)
        if lang == "ig":
            return FileResponse(output_path, media_type="text/plain", filename="response.txt")
        else:
            return FileResponse(output_path, media_type="audio/wav", filename="response.wav")
    return {"transcription": transcription, "ai_response": ai_response}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))