import gradio as gr
import requests
import mimetypes
import os
import threading
import uvicorn
import re

def sanitize(text):
    if isinstance(text, str):
        text = re.sub(r'[\uD800-\uDFFF]', '', text)
        return text.encode("utf-8", "ignore").decode("utf-8")
    return text

def run_fastapi():
    uvicorn.run("main:app", host="0.0.0.0", port=8000)

threading.Thread(target=run_fastapi, daemon=True).start()

API_URLS = {
    "file": "http://127.0.0.1:8000/chat-with-file",
    "url": "http://127.0.0.1:8000/chat-with-url",
    "image": "http://127.0.0.1:8000/extract-text-from-image",
    "audio": "http://127.0.0.1:8000/transcribe-audio"
}

chat_sessions = {"file": [], "url": []}

def format_chat(history):
    return "\n\n".join([f"👤 {sanitize(q)}\n🤖 {sanitize(a)}" for q, a in history])

def ask_from_file(file_obj, question, session_id):
    if not file_obj or not question.strip():
        return "⚠️ Upload a file and enter a question.", chat_sessions["file"]
    if os.path.getsize(file_obj.name) > 10 * 1024 * 1024:
        return "❌ File too large. Max 10MB.", chat_sessions["file"]
    try:
        mime, _ = mimetypes.guess_type(file_obj.name)
        with open(file_obj.name, "rb") as f:
            files = {"file": (file_obj.name, f, mime or "application/octet-stream")}
            history = chat_sessions["file"]
            context = "\n".join([f"Q: {q}\nA: {a}" for q, a in history]) + f"\nQ: {question}"
            res = requests.post(API_URLS["file"], files=files, data={"question": context})
            answer = sanitize(res.json().get("answer", "⚠️ No answer."))
            chat_sessions["file"].append((question, answer))
            return format_chat(chat_sessions["file"]), chat_sessions["file"]
    except Exception as e:
        return f"❌ Error: {e}", chat_sessions["file"]

def ask_from_url(url, question, session_id):
    if not url.strip() or not question.strip():
        return "⚠️ URL and question required.", chat_sessions["url"]
    try:
        history = chat_sessions["url"]
        context = "\n".join([f"Q: {q}\nA: {a}" for q, a in history]) + f"\nQ: {question}"
        res = requests.post(API_URLS["url"], json={"url": url, "question": context})
        answer = sanitize(res.json().get("answer", "⚠️ No answer."))
        chat_sessions["url"].append((question, answer))
        return format_chat(chat_sessions["url"]), chat_sessions["url"]
    except Exception as e:
        return f"❌ Error: {e}", chat_sessions["url"]

def extract_text_from_image(image_path):
    try:
        with open(image_path, "rb") as f:
            files = {"file": ("image.png", f, "image/png")}
            res = requests.post(API_URLS["image"], files=files)
            return sanitize(res.json().get("answer", "⚠️ No text extracted."))
    except Exception as e:
        return f"❌ Error: {e}"

def transcribe_audio(audio_path):
    try:
        with open(audio_path, "rb") as f:
            files = {"file": ("audio.wav", f, "audio/wav")}
            res = requests.post(API_URLS["audio"], files=files)
            return sanitize(res.json().get("answer", "⚠️ No transcript returned."))
    except Exception as e:
        return f"❌ Error: {e}"

def clear_file_chat(): chat_sessions["file"] = []; return "", chat_sessions["file"]
def clear_url_chat(): chat_sessions["url"] = []; return "", chat_sessions["url"]

with gr.Blocks(css="") as demo:
    gr.Markdown("# 🤖 AI Chatbot with File, Web, Image, Audio")
    with gr.Tab("📂 File"):
        file = gr.File(file_types=[".txt", ".csv", ".docx", ".pdf"])
        q1 = gr.Textbox(label="Question")
        b1 = gr.Button("Ask from File")
        out1 = gr.Textbox(label="Chat", lines=10)
        clr1 = gr.Button("Clear")
        b1.click(fn=ask_from_file, inputs=[file, q1, gr.State("file")], outputs=[out1, gr.State("file")])
        clr1.click(fn=clear_file_chat, outputs=[out1, gr.State("file")])

    with gr.Tab("🌐 URL"):
        url = gr.Textbox(label="Website URL")
        q2 = gr.Textbox(label="Question")
        b2 = gr.Button("Ask from URL")
        out2 = gr.Textbox(label="Chat", lines=10)
        clr2 = gr.Button("Clear")
        b2.click(fn=ask_from_url, inputs=[url, q2, gr.State("url")], outputs=[out2, gr.State("url")])
        clr2.click(fn=clear_url_chat, outputs=[out2, gr.State("url")])

    with gr.Tab("🖼️ Image OCR"):
        img = gr.Image(type="filepath")
        b3 = gr.Button("Extract Text")
        out3 = gr.Textbox(label="Text", lines=8)
        b3.click(fn=extract_text_from_image, inputs=img, outputs=out3)

    with gr.Tab("🎤 Audio"):
        aud = gr.Audio(type="filepath")
        b4 = gr.Button("Transcribe")
        out4 = gr.Textbox(label="Transcript", lines=8)
        b4.click(fn=transcribe_audio, inputs=aud, outputs=out4)

demo.launch(share=False)