import os
import time
import asyncio
import shutil

from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

import whisper
from gtts import gTTS

from dotenv import load_dotenv
from openai import OpenAI

# =====================================================
# ENV & OPENAI CLIENT
# =====================================================
load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None

if client:
    print("✅ OpenAI client initialized")
else:
    print("❌ OPENAI_API_KEY not found")

# =====================================================
# FASTAPI APP
# =====================================================
app = FastAPI(title="HF VoiceBot – Whisper + LLM + gTTS")

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

# =====================================================
# TEMP DIRECTORY (HF SAFE)
# =====================================================
TEMP_DIR = "temp"
os.makedirs(TEMP_DIR, exist_ok=True)
app.mount("/api/audio", StaticFiles(directory=TEMP_DIR), name="audio")

# =====================================================
# WHISPER (LAZY LOAD – IMPORTANT FOR HF)
# =====================================================
_whisper_model = None

def get_whisper_model():
    global _whisper_model
    if _whisper_model is None:
        print("⏳ Loading Whisper tiny model (CPU)...")
        _whisper_model = whisper.load_model("tiny")
        print("✅ Whisper model loaded")
    return _whisper_model

# =====================================================
# LANGUAGE MAP (FOR LLM PROMPT)
# =====================================================
LANG_NAMES = {
    "en": "English",
    "hi": "Hindi",
    "te": "Telugu",
    "es": "Spanish",
    "fr": "French",
}

# =====================================================
# MAIN API
# =====================================================
@app.post("/api/process-audio")
async def process_audio(
    audio: UploadFile = File(...),
    input_language: str = Form("en"),
    output_language: str = Form("en"),
):
    if not audio:
        raise HTTPException(status_code=400, detail="No audio uploaded")

    timestamp = int(time.time())
    input_path = f"{TEMP_DIR}/input_{timestamp}.wav"
    output_path = f"{TEMP_DIR}/output_{timestamp}.mp3"

    try:
        # -----------------------------
        # SAVE AUDIO
        # -----------------------------
        with open(input_path, "wb") as f:
            shutil.copyfileobj(audio.file, f)

        # -----------------------------
        # STT (WHISPER)
        # -----------------------------
        model = get_whisper_model()
        loop = asyncio.get_running_loop()

        result = await loop.run_in_executor(
            None,
            lambda: model.transcribe(
                input_path,
                language=input_language,
                fp16=False
            )
        )

        user_text = result.get("text", "").strip()
        print("🎤 USER SAID >>>", user_text)

        # -----------------------------
        # LLM (MANDATORY)
        # -----------------------------
        if not user_text:
            final_response = "Sorry, I could not hear anything."
        elif not client:
            final_response = "LLM is not configured."
        else:
            target_lang = LANG_NAMES.get(output_language, "English")

            response = client.responses.create(
                model="gpt-4o-mini",
                input=f"""
You are a helpful and friendly voice assistant.
Reply ONLY in {target_lang}.
Keep responses short and natural for voice.

User: {user_text}
"""
            )

            final_response = response.output_text.strip()

        print("🤖 FINAL RESPONSE >>>", final_response)

        # -----------------------------
        # TTS (gTTS – HF SUPPORTED)
        # -----------------------------
        tts = gTTS(text=final_response, lang=output_language)
        tts.save(output_path)

        return JSONResponse({
            "user_text": user_text,
            "bot_text": final_response,
            "audio_url": f"/api/audio/{os.path.basename(output_path)}"
        })

    except Exception as e:
        print("❌ ERROR >>>", e)
        return JSONResponse(status_code=500, content={"error": str(e)})

    finally:
        if os.path.exists(input_path):
            os.remove(input_path)

# =====================================================
# HF REQUIRED ENTRYPOINT
# =====================================================
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)