from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import FileResponse
from ultralytics import YOLO
from openai import OpenAI
from gtts import gTTS

import os
import uuid
import requests

app = FastAPI(
    title="Egyptian Talking Artifact API",
    version="1.0"
)

# =====================================
# Folders
# =====================================

os.makedirs("uploads", exist_ok=True)
os.makedirs("outputs", exist_ok=True)

# =====================================
# Models
# =====================================

model = YOLO("best_egypt.pt")

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ["OPENROUTER_API_KEY"]
)

# =====================================
# Wav2Lip API
# =====================================

WAV2LIP_API = "https://fatma812-wav2lip-api.hf.space/generate-video"

# =====================================
# Home
# =====================================

@app.get("/")
def home():
    return {
        "message": "Egyptian Talking Artifact API",
        "docs": "/docs"
    }

# =====================================
# Generate Story + Audio
# =====================================

@app.post("/generate")
async def generate(
    image: UploadFile = File(...),
    language: str = Form(...)
):

    image_name = f"{uuid.uuid4().hex}.jpg"
    image_path = os.path.join("uploads", image_name)

    with open(image_path, "wb") as f:
        f.write(await image.read())

    results = model(image_path)

    if len(results[0].boxes) == 0:
        raise HTTPException(
            status_code=400,
            detail="No artifact detected"
        )

    cls = int(results[0].boxes.cls[0])
    artifact = results[0].names[cls]

    if language == "Arabic":

        prompt = f"""
أنت {artifact}، أثر مصري قديم.

تحدث بصيغة المتكلم.

عرف بنفسك لزوار المتحف.

اذكر معلومة تاريخية قصيرة.

اجعل الكلام مشوقًا.
"""

        tts_lang = "ar"

    else:

        prompt = f"""
You are {artifact}, an ancient Egyptian artifact.

Speak in first person.

Introduce yourself.

Mention one historical fact.

Keep it short.
"""

        tts_lang = "en"

    response = client.chat.completions.create(
        model="openai/gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ]
    )

    story = response.choices[0].message.content

    audio_name = f"{uuid.uuid4().hex}.mp3"
    audio_path = os.path.join("outputs", audio_name)

    gTTS(
        text=story,
        lang=tts_lang
    ).save(audio_path)

    return {
        "artifact": artifact,
        "story": story,
        "audio_url": f"/audio/{audio_name}",
        "image_url": f"/image/{image_name}"
    }

# =====================================
# Download Audio
# =====================================

@app.get("/audio/{filename}")
def get_audio(filename: str):

    path = os.path.join("outputs", filename)

    if not os.path.exists(path):
        raise HTTPException(404, "Audio not found")

    return FileResponse(
        path,
        media_type="audio/mpeg",
        filename=filename
    )

# =====================================
# Download Image
# =====================================

@app.get("/image/{filename}")
def get_image(filename: str):

    path = os.path.join("uploads", filename)

    if not os.path.exists(path):
        raise HTTPException(404, "Image not found")

    return FileResponse(path)

# =====================================
# Generate Talking Video
# =====================================

@app.post("/generate-video")
async def generate_video(

    image: UploadFile = File(...),
    audio: UploadFile = File(...)
):

    response = requests.post(

        WAV2LIP_API,

        files={
            "image": (
                image.filename,
                await image.read(),
                image.content_type
            ),
            "audio": (
                audio.filename,
                await audio.read(),
                audio.content_type
            )
        },

        timeout=600
    )

    if response.status_code != 200:
        raise HTTPException(
            status_code=500,
            detail=response.text
        )

    video_name = f"{uuid.uuid4().hex}.mp4"
    video_path = os.path.join("outputs", video_name)

    with open(video_path, "wb") as f:
        f.write(response.content)

    return FileResponse(
        video_path,
        media_type="video/mp4",
        filename=video_name
    )