import gradio as gr
import torch
import os
import numpy as np
from groq import Groq
import spaces
from transformers import AutoModel, AutoTokenizer
from diffusers import StableDiffusion3Pipeline
from parler_tts import ParlerTTSForConditionalGeneration
import soundfile as sf
from langchain.agents import AgentExecutor, create_react_agent
from langchain.tools import BaseTool
from langchain_groq import ChatGroq
from PIL import Image
from tavily import TavilyClient
import requests
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from langchain.schema import AIMessage
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import RetrievalQA

# Initialize models and clients
MODEL = 'llama3-groq-70b-8192-tool-use-preview'
client = Groq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY"))

vqa_model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_code=True,
                                      device_map="auto", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_code=True)

tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1")
tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-large-v1")

# Updated Image generation model
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Tavily Client for web search
tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API"))

# Function to play voice output
def play_voice_output(response):
    description = "Jon's voice is monotone yet slightly fast in delivery, with a very close recording that almost has no background noise."
    input_ids = tts_tokenizer(description, return_tensors="pt").input_ids.to('cuda')
    prompt_input_ids = tts_tokenizer(response, return_tensors="pt").input_ids.to('cuda')
    generation = tts_model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids)
    audio_arr = generation.cpu().numpy().squeeze()
    sf.write("output.wav", audio_arr, tts_model.config.sampling_rate)
    return "output.wav"

# NumPy Code Calculator Tool
class NumpyCodeCalculator(BaseTool):
    name = "Numpy"
    description = "Useful for performing numpy computations"

    def _run(self, query: str) -> str:
        try:
            local_dict = {"np": np}
            exec(query, local_dict)
            result = local_dict.get("result", "No result found")
            return str(result)
        except Exception as e:
            return f"Error: {e}"

# Web Search Tool
class WebSearch(BaseTool):
    name = "Web"
    description = "Useful for searching the web for information"

    def _run(self, query: str) -> str:
        answer = tavily_client.qna_search(query=query)
        return answer

# Image Generation Tool
class ImageGeneration(BaseTool):
    name = "Image"
    description = "Useful for generating images based on text descriptions"

    def _run(self, query: str) -> str:
        image = pipe(
            query,
            negative_prompt="",
            num_inference_steps=15,
            guidance_scale=7.0,
        ).images[0]
        image.save("output.jpg")
        return "output.jpg"

# Document Question Answering Tool
class DocumentQuestionAnswering(BaseTool):
    name = "Document"
    description = "Useful for answering questions about a specific document"

    def __init__(self, document):
        super().__init__()
        self.document = document
        self.qa_chain = self._setup_qa_chain()

    def _setup_qa_chain(self):
        loader = TextLoader(self.document)
        documents = loader.load()
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
        texts = text_splitter.split_documents(documents)
        embeddings = HuggingFaceEmbeddings()
        db = FAISS.from_documents(texts, embeddings)
        retriever = db.as_retriever()
        qa_chain = RetrievalQA.from_chain_type(
            llm=ChatGroq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY")),
            chain_type="stuff",
            retriever=retriever,
        )
        return qa_chain

    def _run(self, query: str) -> str:
        response = self.qa_chain.run(query)
        return str(response)

# Function to handle different input types and choose the right tool
def handle_input(user_prompt, image=None, audio=None, websearch=False, document=None):
    if audio:
        if isinstance(audio, str):
            audio = open(audio, "rb")
        transcription = client.audio.transcriptions.create(
            file=(audio.name, audio.read()),
            model="whisper-large-v3"
        )
        user_prompt = transcription.text

    tools = [
        NumpyCodeCalculator(),
        ImageGeneration(),
    ]

    # Add the web search tool only if websearch mode is enabled
    if websearch:
        tools.append(WebSearch())

    # Add the document question answering tool only if a document is provided
    if document:
        tools.append(DocumentQuestionAnswering(document))

    llm = ChatGroq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY"))
    agent = create_react_agent(llm, tools, verbose=True)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    if image:
        image = Image.open(image).convert('RGB')
        messages = [{"role": "user", "content": [image, user_prompt]}]
        response = vqa_model.chat(image=None, msgs=messages, tokenizer=tokenizer)
    else:
        response = agent_executor.run(user_prompt)

    return response

    
def create_ui():
    with gr.Blocks(css="""
        /* Overall Styling */
        body {
            font-family: 'Poppins', sans-serif;
            background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
            margin: 0;
            padding: 0;
            color: #333;
        }

        /* Title Styling */
        .gradio-container h1 {
            text-align: center;
            padding: 20px 0;
            background: linear-gradient(45deg, #007bff, #00c6ff);
            color: white;
            font-size: 2.5em;
            font-weight: bold;
            letter-spacing: 1px;
            text-transform: uppercase;
            margin: 0;
            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.2);
        }

        /* Input Area Styling */
        .gradio-container .gr-row {
            display: flex;
            justify-content: space-around;
            align-items: center;
            padding: 20px;
            background-color: white;
            border-radius: 10px;
            box-shadow: 0px 6px 12px rgba(0, 0, 0, 0.1);
            margin-bottom: 20px;
        }

        .gradio-container .gr-column {
            flex: 1;
            margin: 0 10px;
        }

        /* Textbox Styling */
        .gradio-container textarea {
            width: calc(100% - 20px);
            padding: 15px;
            border: 2px solid #007bff;
            border-radius: 8px;
            font-size: 1.1em;
            transition: border-color 0.3s, box-shadow 0.3s;
        }

        .gradio-container textarea:focus {
            border-color: #00c6ff;
            box-shadow: 0px 0px 8px rgba(0, 198, 255, 0.5);
            outline: none;
        }

        /* Button Styling */
        .gradio-container button {
            background: linear-gradient(45deg, #007bff, #00c6ff);
            color: white;
            padding: 15px 25px;
            border: none;
            border-radius: 8px;
            cursor: pointer;
            font-size: 1.2em;
            font-weight: bold;
            transition: background 0.3s, transform 0.3s;
            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);
        }

        .gradio-container button:hover {
            background: linear-gradient(45deg, #0056b3, #009bff);
            transform: translateY(-3px);
        }

        .gradio-container button:active {
            transform: translateY(0);
        }

        /* Output Area Styling */
        .gradio-container .output-area {
            padding: 20px;
            text-align: center;
            background-color: #f7f9fc;
            border-radius: 10px;
            box-shadow: 0px 6px 12px rgba(0, 0, 0, 0.1);
            margin-top: 20px;
        }

        /* Image Styling */
        .gradio-container img {
            max-width: 100%;
            height: auto;
            border-radius: 10px;
            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);
            transition: transform 0.3s, box-shadow 0.3s;
        }

        .gradio-container img:hover {
            transform: scale(1.05);
            box-shadow: 0px 6px 12px rgba(0, 0, 0, 0.2);
        }

        /* Checkbox Styling */
        .gradio-container input[type="checkbox"] {
            width: 20px;
            height: 20px;
            cursor: pointer;
            accent-color: #007bff;
            transition: transform 0.3s;
        }

        .gradio-container input[type="checkbox"]:checked {
            transform: scale(1.2);
        }

        /* Audio and Document Upload Styling */
        .gradio-container .gr-file-upload input[type="file"] {
            width: 100%;
            padding: 10px;
            border: 2px solid #007bff;
            border-radius: 8px;
            cursor: pointer;
            background-color: white;
            transition: border-color 0.3s, background-color 0.3s;
        }

        .gradio-container .gr-file-upload input[type="file"]:hover {
            border-color: #00c6ff;
            background-color: #f0f8ff;
        }

        /* Advanced Tooltip Styling */
        .gradio-container .gr-tooltip {
            position: relative;
            display: inline-block;
            cursor: pointer;
        }

        .gradio-container .gr-tooltip .tooltiptext {
            visibility: hidden;
            width: 200px;
            background-color: black;
            color: #fff;
            text-align: center;
            border-radius: 6px;
            padding: 5px;
            position: absolute;
            z-index: 1;
            bottom: 125%;
            left: 50%;
            margin-left: -100px;
            opacity: 0;
            transition: opacity 0.3s;
        }

        .gradio-container .gr-tooltip:hover .tooltiptext {
            visibility: visible;
            opacity: 1;
        }

        /* Footer Styling */
        .gradio-container footer {
            text-align: center;
            padding: 10px;
            background: #007bff;
            color: white;
            font-size: 0.9em;
            border-radius: 0 0 10px 10px;
            box-shadow: 0px -2px 8px rgba(0, 0, 0, 0.1);
        }

    """) as demo:
        gr.Markdown("# AI Assistant")
        with gr.Row():
            with gr.Column(scale=2):
                user_prompt = gr.Textbox(placeholder="Type your message here...", lines=1)
            with gr.Column(scale=1):
                image_input = gr.Image(type="filepath", label="Upload an image", elem_id="image-icon")
                audio_input = gr.Audio(type="filepath", label="Upload audio", elem_id="mic-icon")
                document_input = gr.File(type="filepath", label="Upload a document", elem_id="document-icon")
                voice_only_mode = gr.Checkbox(label="Enable Voice Only Mode", elem_id="voice-only-mode")
                websearch_mode = gr.Checkbox(label="Enable Web Search", elem_id="websearch-mode")
            with gr.Column(scale=1):
                submit = gr.Button("Submit")

        output_label = gr.Label(label="Output")
        audio_output = gr.Audio(label="Audio Output", visible=False)

        submit.click(
            fn=main_interface,
            inputs=[user_prompt, image_input, audio_input, voice_only_mode, websearch_mode, document_input],
            outputs=[output_label, audio_output]
        )

        voice_only_mode.change(
            lambda x: gr.update(visible=not x),
            inputs=voice_only_mode,
            outputs=[user_prompt, image_input, websearch_mode, document_input, submit]
        )
        voice_only_mode.change(
            lambda x: gr.update(visible=x),
            inputs=voice_only_mode,
            outputs=[audio_input]
        )

    return demo

# Main interface function
@spaces.GPU()
def main_interface(user_prompt, image=None, audio=None, voice_only=False, websearch=False, document=None):
    print("Starting main_interface function")
    vqa_model.to(device='cuda', dtype=torch.bfloat16)
    tts_model.to("cuda")
    pipe.to("cuda")

    print(f"user_prompt: {user_prompt}, image: {image}, audio: {audio}, voice_only: {voice_only}, websearch: {websearch}, document: {document}")

    try:
        response = handle_input(user_prompt, image=image, audio=audio, websearch=websearch, document=document)
        print("handle_input function executed successfully")
    except Exception as e:
        print(f"Error in handle_input: {e}")
        response = "Error occurred during processing."

    if voice_only:
        try:
            audio_output = play_voice_output(response)
            print("play_voice_output function executed successfully")
            return "Response generated.", audio_output
        except Exception as e:
            print(f"Error in play_voice_output: {e}")
            return "Error occurred during voice output.", None
    else:
        return response, None

# Launch the UI
demo = create_ui()
demo.launch()