import os
import gradio as gr
import openai

# Optional imports for extended functionality
import base64
import fitz  # PyMuPDF
from io import BytesIO
from moviepy.video.io.VideoFileClip import VideoFileClip
from gtts import gTTS
from gradio_client import Client

# 🔐 CONFIG
KEY_FILE = "openai_api_key.txt"
MODEL = "gpt-4o-2024-05-13"

# 🔧 Helpers

def save_api_key(key: str) -> str:
    """Save a new OpenAI API key to disk."""
    with open(KEY_FILE, "w") as f:
        f.write(key.strip())
    return "🔑 Key saved!"


def get_api_key(input_key: str) -> str:
    """Determine which OpenAI key to use: input > saved > env."""
    # 1️⃣ User-provided key
    if input_key and input_key.strip():
        save_api_key(input_key)
        return input_key.strip()

    # 2️⃣ Previously saved key
    if os.path.exists(KEY_FILE):
        with open(KEY_FILE) as f:
            return f.read().strip()

    # 3️⃣ Environment variable
    env_key = os.getenv("OPENAI_KEY", "")
    if env_key:
        save_api_key(env_key)
        return env_key

    # ❗ No key found
    raise gr.Error("❗ OpenAI API key required.")


def apply_filters(files, filters):
    """
    Filter uploaded files by extension and user-selected filters.
    Returns a list for the Gallery display and tracking info for table generation.
    """
    gallery_items = []
    tracking = []

    if not files:
        return gallery_items, tracking

    for f in files:
        path = getattr(f, "name", f)
        ext = os.path.splitext(path)[1].lower()

        if ext in ['.png', '.jpg', '.jpeg'] and '🖼️ Images' in filters:
            gallery_items.append(f)
            tracking.append((path, 'image'))
        if ext in ['.wav', '.mp3'] and '🎤 Audio' in filters:
            gallery_items.append(f)
            tracking.append((path, 'audio'))
        if ext in ['.mp4', '.webm'] and '🎥 Video' in filters:
            gallery_items.append(f)
            tracking.append((path, 'video'))
        if ext == '.pdf' and '📄 PDF' in filters:
            gallery_items.append(f)
            tracking.append((path, 'pdf'))

    return gallery_items, tracking


def generate_table(selected, api_key):
    """
    Build a Markdown table of selected items with type icons.
    """
    key = get_api_key(api_key)
    openai.api_key = key

    md = "| ✅ | Type | Filename |\n|---|------|----------|\n"
    emoji_map = {
        'image': '🖼️',
        'audio': '🎤',
        'video': '🎥',
        'pdf': '📄'
    }

    for path, typ in selected:
        name = os.path.basename(path)
        md += f"| ✅ | {emoji_map.get(typ, '')} {typ.capitalize()} | {name} |\n"

    return md


def chat_handler(api_key, message, history):
    """
    Send a user message + chat history to OpenAI and append the response.
    """
    key = get_api_key(api_key)
    openai.api_key = key

    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": message})
    resp = openai.ChatCompletion.create(model=MODEL, messages=messages)

    answer = resp.choices[0].message.content
    history.append((message, answer))
    return history

# 🔑 UI Definition
with gr.Blocks(title="🔬🧠 ScienceBrain.Gradio") as demo:
    gr.Markdown("# 🔬🧠 ScienceBrain Gradio")

    # --- API Key Input & Save ---
    with gr.Row():
        api_input = gr.Textbox(
            label="🔑 OpenAI Key", type="password", value=os.getenv("OPENAI_KEY", "")
        )
        save_btn = gr.Button("💾 Save Key")
        status_txt = gr.Textbox(interactive=False)
        save_btn.click(save_api_key, inputs=api_input, outputs=status_txt)

    # --- Media Gallery & Filters ---
    gr.Markdown("## 📋 Media Gallery & Filters")
    upload = gr.File(file_count="multiple", label="Upload files (images, audio, videos, PDFs)")
    gallery = gr.Gallery(label="Filtered Gallery", columns=4)
    filter_opts = ["🖼️ Images", "🎤 Audio", "🎥 Video", "📄 PDF"]
    filters = gr.CheckboxGroup(filter_opts, value=filter_opts, label="🔍 Filter Types")
    select_btn = gr.Button("⚙️ Apply Filters")
    selected = gr.State()

    select_btn.click(
        apply_filters,
        inputs=[upload, filters],
        outputs=[gallery, selected]
    )

    # --- Summary Table ---
    gr.Markdown("## 📝 Summary Table")
    table_md = gr.Markdown()
    table_btn = gr.Button("Generate Table")
    table_btn.click(
        generate_table,
        inputs=[selected, api_input],
        outputs=[table_md]
    )

    # --- Chat Tab ---
    with gr.Tab("💬 Chat"):
        chatbot = gr.Chatbot()
        msg_in = gr.Textbox(placeholder="Type your message...")
        msg_in.submit(
            chat_handler,
            inputs=[api_input, msg_in, chatbot],
            outputs=[chatbot]
        )

if __name__ == "__main__":
    demo.launch()