import gradio as gr
import os
import uuid
import requests

from ultralytics import YOLO
from openai import OpenAI
from gtts import gTTS

# =========================
# YOLO
# =========================

model = YOLO("best_egypt.pt")

# =========================
# OpenRouter
# =========================

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

# =========================
# FastAPI URL
# =========================

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

# =========================
# Story + Voice
# =========================

def generate(image_path, language):

    results = model(image_path)

    if len(results[0].boxes) == 0:
        raise gr.Error("No artifact detected")

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

    if language == "Arabic":

        prompt = f"""
أنت {artifact} من آثار مصر القديمة.

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

احك قصة قصيرة لا تزيد عن 3 جمل.
"""

        tts_lang = "ar"

    else:

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

Speak in first person.

Tell a short interesting story in 3 sentences.
"""

        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_path = f"audio_{uuid.uuid4().hex}.mp3"

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

    return artifact, story, audio_path


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

def make_video(image_path, audio_path):

    if image_path is None:
        raise gr.Error("Upload image first")

    if audio_path is None:
        raise gr.Error("Generate audio first")

    with open(image_path, "rb") as img, open(audio_path, "rb") as aud:

        response = requests.post(
            WAV2LIP_API,
            files={
                "image": img,
                "audio": aud
            },
            timeout=600
        )

    if response.status_code != 200:
        raise gr.Error(response.text)

    video_path = f"video_{uuid.uuid4().hex}.mp4"

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

    return video_path


# =========================
# UI
# =========================

with gr.Blocks() as demo:

    gr.Markdown("# 🏛 Talking Egyptian Artifact AI")

    image = gr.Image(
        type="filepath",
        label="Artifact Image"
    )

    language = gr.Radio(
        ["Arabic", "English"],
        value="Arabic",
        label="Language"
    )

    btn_generate = gr.Button("Generate Story + Voice")

    artifact = gr.Textbox(label="Artifact")

    story = gr.Textbox(
        label="Story",
        lines=6
    )

    audio = gr.Audio(
        type="filepath",
        label="Generated Voice"
    )

    btn_video = gr.Button("Generate Talking Video")

    video = gr.Video(
        label="Talking Video"
    )

    btn_generate.click(
        fn=generate,
        inputs=[image, language],
        outputs=[artifact, story, audio]
    )

    btn_video.click(
        fn=make_video,
        inputs=[image, audio],
        outputs=video
    )

demo.queue()
demo.launch()