import os
import requests
import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv
from helpers import format_images, format_videos, format_links, format_local_map, format_knowledge

# Load API key from .env
load_dotenv()
AIML_API_KEY = os.getenv('AIML_API_KEY')
API_URL = 'https://api.aimlapi.com'

# Bagoodex API client encapsulation
class BagoodexClient:
    def __init__(self, api_key=AIML_API_KEY, api_url=API_URL):
        self.api_key = api_key
        self.api_url = api_url
        self.client = OpenAI(base_url=self.api_url, api_key=self.api_key)

    def complete_chat(self, query):
        response = self.client.chat.completions.create(
            model="bagoodex/bagoodex-search-v1",
            messages=[{"role": "user", "content": query}],
        )
        followup_id = response.id  # unique ID for followup searches
        answer = response.choices[0].message.content
        return followup_id, answer

    def get_links(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(f"{self.api_url}/v1/bagoodex/links", headers=headers, params=params)
        return response.json()

    def get_images(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(f"{self.api_url}/v1/bagoodex/images", headers=headers, params=params)
        return response.json()

    def get_videos(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(f"{self.api_url}/v1/bagoodex/videos", headers=headers, params=params)
        return response.json()

    def get_local_map(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(f"{self.api_url}/v1/bagoodex/local-map", headers=headers, params=params)
        return response.json()

    def get_knowledge(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(f"{self.api_url}/v1/bagoodex/knowledge", headers=headers, params=params)
        return response.json()

client = BagoodexClient()

# Callback functions
def chat_response(user_input, chat_history, followup_id):
    followup_id_new, answer = client.complete_chat(user_input)
    chat_history = chat_history or []
    chat_history.append((user_input, answer))
    return "", chat_history, followup_id_new

def perform_local_map_search(followup_id, chat_history):
    if not followup_id:
        return chat_history
    result = client.get_local_map(followup_id)
    formatted = format_local_map(result)
    chat_history.append((None, formatted))
    return chat_history

def perform_knowledge_search(followup_id, chat_history):
    if not followup_id:
        return chat_history
    result = client.get_knowledge(followup_id)
    formatted = format_knowledge(result)
    chat_history.append((None, formatted))
    return chat_history

def perform_image_search(followup_id):
    if not followup_id:
        return []
    result = client.get_images(followup_id)
    urls = format_images(result)
    return urls

def perform_video_search(followup_id):
    if not followup_id:
        return "No followup ID available."
    result = client.get_videos(followup_id)
    return format_videos(result)

def perform_links_search(followup_id):
    if not followup_id:
        return "No followup ID available."
    result = client.get_links(followup_id)
    return format_links(result)

# Custom CSS to help the chat component fill the available height
css = """
#chatbot {
    height: 100%;
}
"""

# Build UI
with gr.Blocks(css=css, fill_height=True) as demo:
    with gr.Row():
        # Left column: Chat interface
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="Bagoodex Chat", elem_id="chatbot")
            # State variables: chat history and followup ID.
            chat_history_state = gr.State([])
            followup_state = gr.State(None)
            # Two small buttons placed above the input field:
            with gr.Row():
                btn_local_map = gr.Button("Local Map Search", variant="secondary", size="sm")
                btn_knowledge = gr.Button("Knowledge Base", variant="secondary", size="sm")
            # Chat input area: textbox and explicit submit button.
            with gr.Row():
                txt = gr.Textbox(placeholder="Enter your query here...", show_label=False, scale=7)
                submit_btn = gr.Button("Submit", scale=1)
            # Wire up the chat callbacks:
            submit_btn.click(chat_response, inputs=[txt, chat_history_state, followup_state],
                               outputs=[txt, chatbot, followup_state])
            txt.submit(chat_response, inputs=[txt, chat_history_state, followup_state],
                       outputs=[txt, chatbot, followup_state])
            # The two additional search buttons update the chat history by appending a new assistant message.
            btn_local_map.click(perform_local_map_search, inputs=[followup_state, chat_history_state],
                                outputs=[chatbot, chat_history_state])
            btn_knowledge.click(perform_knowledge_search, inputs=[followup_state, chat_history_state],
                                outputs=[chatbot, chat_history_state])
        # Right column: Advanced search options for images, videos, and links.
        with gr.Column(scale=1):
            gr.Markdown("### Advanced Search Options")
            with gr.Column(variant="panel"):
                btn_images = gr.Button("Search Images")
                btn_videos = gr.Button("Search Videos")
                btn_links = gr.Button("Search Links")
                gallery_output = gr.Gallery(label="Image Results", columns=2)
                video_output = gr.HTML(label="Video Results")
                links_output = gr.HTML(label="Links Results")
                btn_images.click(perform_image_search, inputs=[followup_state], outputs=[gallery_output])
                btn_videos.click(perform_video_search, inputs=[followup_state], outputs=[video_output])
                btn_links.click(perform_links_search, inputs=[followup_state], outputs=[links_output])
    demo.launch()