import gradio as gr
import httpx
import json

API_URL = "https://maleshkumar12345--nexus-multimodal-backend-serve.modal.run/api"

async def upload_invoice(image_filepath):
    if not image_filepath:
        return "Please upload an image.", None
    
    async with httpx.AsyncClient(follow_redirects=True) as client:
        with open(image_filepath, "rb") as f:
            files = {"file": (image_filepath, f, "image/jpeg")}
            try:
                response = await client.post(f"{API_URL}/upload-invoice", files=files, timeout=60.0)
                response.raise_for_status()
                data = response.json()
                doc_id = data["id"]
                extracted = json.dumps(data["data"].get("extracted_data", {}), indent=2)
                return extracted, doc_id
            except Exception as e:
                return f"Error: {str(e)}", None

async def upload_video(video_filepath):
    if not video_filepath:
        return "Please upload a video.", None
    
    async with httpx.AsyncClient(follow_redirects=True) as client:
        with open(video_filepath, "rb") as f:
            files = {"file": (video_filepath, f, "video/mp4")}
            try:
                response = await client.post(f"{API_URL}/upload-video", files=files, timeout=300.0)
                response.raise_for_status()
                data = response.json()
                doc_id = data["id"]
                insights = data["data"].get("insights", {})
                return json.dumps(insights, indent=2), doc_id
            except Exception as e:
                return f"Error: {str(e)}", None

async def chat(message, history, doc_id, doc_type):
    if not doc_id:
        history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": f"Please upload a {doc_type} first."}
        ]
        return "", history
    
    # Build chat history for the API payload
    chat_history = []
    for msg in history:
        chat_history.append({"role": msg["role"], "content": msg["content"]})
            
    payload = {
        "document_id": doc_id,
        "document_type": doc_type,
        "message": message,
        "history": chat_history
    }
    
    async with httpx.AsyncClient(follow_redirects=True) as client:
        try:
            response = await client.post(f"{API_URL}/chat", json=payload, timeout=60.0)
            response.raise_for_status()
            reply = response.json().get("reply", "")
            history = history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": reply}
            ]
            return "", history
        except Exception as e:
            history = history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": f"Error: {str(e)}"}
            ]
            return "", history

with gr.Blocks(title="Multimodal AI Assistant") as demo:
    gr.Markdown("# Multimodal AI Assistant")
    gr.Markdown("Upload an invoice to extract data or a video to extract insights, and then chat with the extracted information.")
    
    with gr.Tabs():
        with gr.Tab("Invoice Analysis"):
            with gr.Row():
                with gr.Column():
                    invoice_image = gr.Image(type="filepath", label="Upload Invoice")
                    invoice_btn = gr.Button("Analyze Invoice", variant="primary")
                    invoice_id = gr.State(value=None)
                with gr.Column():
                    invoice_output = gr.Code(label="Extracted Data", language="json")
            
            with gr.Row():
                invoice_chatbot = gr.Chatbot(label="Chat about Invoice")
            with gr.Row():
                invoice_msg = gr.Textbox(show_label=False, placeholder="Ask a question about the invoice...", scale=4)
                invoice_send = gr.Button("Send", variant="primary", scale=1)
                
            invoice_btn.click(
                upload_invoice,
                inputs=[invoice_image],
                outputs=[invoice_output, invoice_id]
            )
            invoice_send.click(
                chat,
                inputs=[invoice_msg, invoice_chatbot, invoice_id, gr.State("invoice")],
                outputs=[invoice_msg, invoice_chatbot]
            )
            invoice_msg.submit(
                chat,
                inputs=[invoice_msg, invoice_chatbot, invoice_id, gr.State("invoice")],
                outputs=[invoice_msg, invoice_chatbot]
            )

        with gr.Tab("Video Insights"):
            with gr.Row():
                with gr.Column():
                    video_file = gr.Video(label="Upload Video")
                    video_btn = gr.Button("Analyze Video", variant="primary")
                    video_id = gr.State(value=None)
                with gr.Column():
                    video_output = gr.Code(label="Video Insights", language="json")
            
            with gr.Row():
                video_chatbot = gr.Chatbot(label="Chat about Video")
            with gr.Row():
                video_msg = gr.Textbox(show_label=False, placeholder="Ask a question about the video...", scale=4)
                video_send = gr.Button("Send", variant="primary", scale=1)
                
            video_btn.click(
                upload_video,
                inputs=[video_file],
                outputs=[video_output, video_id]
            )
            video_send.click(
                chat,
                inputs=[video_msg, video_chatbot, video_id, gr.State("video")],
                outputs=[video_msg, video_chatbot]
            )
            video_msg.submit(
                chat,
                inputs=[video_msg, video_chatbot, video_id, gr.State("video")],
                outputs=[video_msg, video_chatbot]
            )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())