import gradio as gr
import os
import requests
from huggingface_hub import InferenceClient
import google.generativeai as genai
import openai

def api_check_msg(api_key, selected_model):
    res = validate_api_key(api_key, selected_model)
    return res["message"]

def validate_api_key(api_key, selected_model):
    # Check if the API key is valid for GPT-3.5-Turbo
    if "GPT" in selected_model:
        url = "https://api.openai.com/v1/models"
        headers = {
            "Authorization": f"Bearer {api_key}"
        }
        try:
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
            else:
                return {"is_valid": False, "message": f'<p style="color: red;">Invalid OpenAI API Key. Status code: {response.status_code}</p>'}
        except requests.exceptions.RequestException as e:
            return {"is_valid": False, "message": f'<p style="color: red;">Invalid OpenAI API Key. Error: {e}</p>'}
    elif "Llama" in selected_model:
        url = "https://huggingface.co/api/whoami-v2"
        headers = {
            "Authorization": f"Bearer {api_key}"
        }
        try:
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
            else:
                return {"is_valid": False, "message": f'<p style="color: red;">Invalid Hugging Face API Key. Status code: {response.status_code}</p>'}
        except requests.exceptions.RequestException as e:
            return {"is_valid": False, "message": f'<p style="color: red;">Invalid Hugging Face API Key. Error: {e}</p>'}
    elif "Gemini" in selected_model:
        try:
            genai.configure(api_key=api_key)
            model = genai.GenerativeModel("gemini-1.5-flash")
            response = model.generate_content("Help me diagnose the patient.")
            return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
        except Exception as e:
            return {"is_valid": False, "message": f'<p style="color: red;">Invalid Google API Key. Error: {e}</p>'}

def generate_text_chatgpt(key, prompt, temperature, top_p):

    openai.api_key = key

    prompt_template = f"""
    {prompt} <Choose only one among the words Psoriasis, Arthritis, Bronchial asthma or Cervical spondylosis>
    """

    response = openai.chat.completions.create(
      model="gpt-3.5-turbo-1106",
      messages=[{"role": "system", "content": "You are a talented diagnostician who is diagnosing a patient."},
                {"role": "user", "content": prompt_template}],
      temperature=temperature,
      max_tokens=50,
      top_p=top_p,
      frequency_penalty=0
    )

    return response.choices[0].message.content


def generate_text_gemini(key, prompt, temperature, top_p):
    genai.configure(api_key=key)

    prompt_template = f"""
    {prompt} <Choose only one among the words Psoriasis, Arthritis, Bronchial asthma or Cervical spondylosis>
    """

    generation_config = genai.GenerationConfig(
        max_output_tokens=len(prompt_template)+50,
        temperature=temperature,
        top_p=top_p,
    )
    model = genai.GenerativeModel("gemini-1.5-flash", generation_config=generation_config)
    response = model.generate_content(prompt_template)
    return response.text


def generate_text_llama(key, prompt, temperature, top_p):
    model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
    client = InferenceClient(api_key=key)

    prompt_template = f"""
    {prompt} <Choose only one among the words Psoriasis, Arthritis, Bronchial asthma or Cervical spondylosis>
    Do not list the symptoms again in the response. Do not add any additional text. Do not attempt to explain your answer.
    """

    messages = [{"role": "system", "content": "You are a talented diagnostician who is diagnosing a patient."},
                {"role": "user","content": prompt_template}]

    completion = client.chat.completions.create(
        model=model_name,
        messages=messages, 
        max_tokens=len(prompt_template)+50,
        temperature=temperature,
        top_p=top_p
    )

    response = completion.choices[0].message.content
    if len(response) > len(prompt_template):
        return response[len(prompt_template):]
    return response


def diagnose(key, model, top_k, temperature, symptom_prompt):

    model_map = {
        "GPT-3.5-Turbo": "GPT",
        "Llama-3": "Llama",
        "Gemini-1.5": "Gemini"
    }
    if symptom_prompt:
        if "GPT" in model:
            message = generate_text_chatgpt(key, symptom_prompt, temperature, top_k)
        elif "Llama" in model:
            message = generate_text_llama(key, symptom_prompt, temperature, top_k)
        elif "Gemini" in model:
            message = generate_text_gemini(key, symptom_prompt, temperature, top_k)
        else:
            message = "Incorrect model, please try again."
    else:
        message = "Please add the symptoms data"
    
    return message

def update_model_components(selected_model):
    model_map = {
                "GPT-3.5-Turbo": "GPT",
                "Llama-3": "Llama",
                "Gemini-1.5": "Gemini"
            }

    link_map = {
        "GPT-3.5-Turbo": "https://platform.openai.com/account/api-keys",
        "Llama-3": "https://hf.co/settings/tokens",
        "Gemini-1.5": "https://aistudio.google.com/apikey"
    }
    textbox_label = f"Please input the API key for your {model_map[selected_model]} model"
    button_value = f"Don't have an API key? Get one for the {model_map[selected_model]} model here."
    button_link = link_map[selected_model]
    return gr.update(label=textbox_label), gr.update(value=button_value, link=button_link)

def toggle_button(symptoms_text, api_key, model):
    if symptoms_text.strip() and validate_api_key(api_key, model):
        return gr.update(interactive=True)
    return gr.update(interactive=False)


with gr.Blocks() as ui:

    with gr.Row(equal_height=500):
        with gr.Column(scale=1, min_width=300):
            model = gr.Radio(label="LLM Selection", value="GPT-3.5-Turbo",  
                             choices=["GPT-3.5-Turbo", "Llama-3", "Gemini-1.5"])
            is_valid = False
            key = gr.Textbox(label="Please input the API key for your Large Language model", type="password")
            status_message = gr.HTML(label="Validation Status")
            key.input(fn=api_check_msg, inputs=[key, model], outputs=status_message)
            button = gr.Button(value="Don't have an API key? Get one for the GPT model here.", link="https://platform.openai.com/account/api-keys")
            model.change(update_model_components, inputs=model, outputs=[key, button])
            gr.ClearButton(key, variant="primary")
            
        with gr.Column(scale=2, min_width=600):
            gr.Markdown("## Hello, Welcome to the GUI by Team #9.")
            temperature = gr.Slider(0.0, 1.0, value=0.7, step = 0.05, label="Temperature", info="Set the Temperature")
            top_p = gr.Slider(0.0, 1.0, value=0.9, step = 0.05, label="top-p value", info="Set the sampling nucleus parameter")
            symptoms = gr.Textbox(label="Add the symptom data in the input to receive diagnosis")
            llm_btn = gr.Button(value="Diagnose Disease", variant="primary", elem_id="diagnose", interactive=False)
            symptoms.input(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
            key.input(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
            model.change(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
            output = gr.Textbox(label="LLM Output Status", interactive=False, placeholder="Output will appear here...")
            llm_btn.click(fn=diagnose, inputs=[key, model, top_p, temperature, symptoms], outputs=output, api_name="LLM_Comparator")


ui.launch(share=True)