import gradio as gr
import requests
import time
import json
from typing import Dict, Any, List, Tuple
import os
from dotenv import load_dotenv
from utils.url_config import API_BASE_URL

load_dotenv()

# Get collection name from environment variable
DEFAULT_COLLECTION = "mangoit_docs"
RAG_COLLECTION = os.getenv("RAG_COLLECTION", DEFAULT_COLLECTION)

class MangoITChatbot:
    def __init__(self, api_url: str = API_BASE_URL):
        self.api_url = api_url
        self.conversation_id = None
        self.active_task_id = None
        self.polling = False
        
    def send_message(self, message: str) -> Tuple[str, str]:
        """
        Send a message to the chatbot API
        
        Args:
            message: The user's message
            
        Returns:
            Tuple[str, str]: (task_id, conversation_id)
        """
        url = f"{self.api_url}/rag-chat"
        payload = {
            "message": message,
            "collection_name": RAG_COLLECTION
        }
        
        # Include conversation ID if we have one
        if self.conversation_id:
            payload["conversation_id"] = self.conversation_id
            
        try:
            response = requests.post(url, json=payload)
            data = response.json()
            
            task_id = data.get("task_id")
            self.conversation_id = data.get("conversation_id")
            
            return task_id, self.conversation_id
        except Exception as e:
            print(f"Error sending message: {str(e)}")
            return None, self.conversation_id
            
    def check_status(self, task_id: str) -> Dict[str, Any]:
        """
        Check the status of a task
        
        Args:
            task_id: The ID of the task
            
        Returns:
            Dict[str, Any]: The task status
        """
        if not task_id:
            return {"status": "error", "error": "No task ID provided"}
            
        url = f"{self.api_url}/rag-chat/status/{task_id}"
        
        try:
            response = requests.get(url)
            return response.json()
        except Exception as e:
            print(f"Error checking status: {str(e)}")
            return {"status": "error", "error": str(e)}

# Initialize the chatbot
chatbot = MangoITChatbot()

# Store conversation history
conversation_history = []

def user_message(message: str, history: List[List[str]]) -> Tuple[List[List[str]], str, str]:
    """
    Process a user message
    
    Args:
        message: The user's message
        history: The conversation history
        
    Returns:
        Tuple[List[List[str]], str, str]: (updated_history, status_html, conversation_id)
    """
    # Add user message to history
    history.append([message, None])
    
    # Send message to API
    task_id, conversation_id = chatbot.send_message(message)
    
    if not task_id:
        history[-1][1] = "Error: Could not send message to the server. Is the FastAPI server running?"
        return history, "<span style='color: red'>Error: Server connection failed</span>", ""
    
    # Return immediately with a placeholder
    return history, f"<span style='color: blue'>Processing... Task ID: {task_id}</span>", conversation_id

def bot_response(history: List[List[str]], status_html: str, conversation_id: str) -> Tuple[List[List[str]], str]:
    """
    Get the bot's response by polling the API
    
    Args:
        history: The conversation history
        status_html: The current status HTML
        conversation_id: The conversation ID
        
    Returns:
        Tuple[List[List[str]], str]: (updated_history, updated_status_html)
    """
    # Extract task_id from status_html
    import re
    match = re.search(r"Task ID: ([a-f0-9-]+)", status_html)
    if not match:
        return history, "<span style='color: red'>Error: No task ID found</span>"
        
    task_id = match.group(1)
    
    # Poll for results
    max_attempts = 30
    for attempt in range(max_attempts):
        status = chatbot.check_status(task_id)
        
        # Update status with progress
        progress = status.get("progress", 0)
        status_html = f"<span style='color: blue'>Processing... {progress}% complete</span>"
        
        # Check if completed
        if status.get("status") == "completed" and status.get("result"):
            result = status.get("result", {})
            response = result.get("response", "No response generated")
            
            # Update history with bot response
            history[-1][1] = response
            
            # Format sources if available
            sources = result.get("sources", [])
            if sources:
                source_html = "<br><small><b>Sources:</b><ul>"
                for source in sources:
                    source_html += f"<li>{source.get('source', 'Unknown')}</li>"
                source_html += "</ul></small>"
                status_html = source_html
            else:
                status_html = ""
                
            return history, status_html
            
        # Check for errors
        if status.get("status") in ["failed", "timeout"]:
            error_msg = status.get("error", "Unknown error")
            history[-1][1] = f"Error: {error_msg}"
            return history, f"<span style='color: red'>Error: {error_msg}</span>"
            
        # Wait before polling again
        time.sleep(1)
    
    # If we get here, we timed out
    history[-1][1] = "Sorry, the request timed out. Please try again."
    return history, "<span style='color: red'>Error: Request timed out</span>"

def create_new_conversation() -> Tuple[List[List[str]], str, str]:
    """Reset the conversation"""
    chatbot.conversation_id = None
    return [], "", ""

# Create the Gradio interface
with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown("# MangoIT Chatbot - Sam")
    gr.Markdown("Ask Sam about MangoIT's services, technologies, and expertise.")
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot_ui = gr.Chatbot(
                [],
                elem_id="chatbot",
                height=500,
                avatar_images=("👤", "🤖")
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    show_label=False,
                    placeholder="Ask Sam a question...",
                    container=False
                )
                submit_btn = gr.Button("Send", variant="primary")
        
        with gr.Column(scale=1):
            status_display = gr.HTML("")
            conversation_id_display = gr.Textbox(label="Conversation ID", interactive=False)
            new_chat_btn = gr.Button("New Conversation")
    
    # Set up event handlers
    msg_event = msg.submit(
        user_message, 
        [msg, chatbot_ui], 
        [chatbot_ui, status_display, conversation_id_display],
        queue=False
    ).then(
        bot_response,
        [chatbot_ui, status_display, conversation_id_display],
        [chatbot_ui, status_display]
    )
    
    submit_btn.click(
        user_message, 
        [msg, chatbot_ui], 
        [chatbot_ui, status_display, conversation_id_display],
        queue=False
    ).then(
        bot_response,
        [chatbot_ui, status_display, conversation_id_display],
        [chatbot_ui, status_display]
    )
    
    new_chat_btn.click(
        create_new_conversation,
        [],
        [chatbot_ui, status_display, conversation_id_display]
    )
    
    # Clear the message box after sending
    msg_event.then(lambda: "", None, msg)
    submit_btn.click(lambda: "", None, msg)

# Launch the interface
if __name__ == "__main__":
    print("Starting Gradio interface. Make sure the FastAPI server is running!")
    print("Run 'uvicorn main:app --reload' in a separate terminal if it's not running.")
    demo.launch(share=False)