import json
import gradio as gr
import os
from huggingface_hub import InferenceClient, login
from image_agent import init_stored_image_agent, stored_images
from dynamic_image_agent import main
from mulitagents import define_multi_agent
from nsfw_detection import classify_image_if_nsfw

login(os.environ.get('HF_TOKEN'))
os.environ["OPENAI_API_KEY"] = os.environ.get('OPENAI_API_KEY')

# --- Global state ---
active_agent = None  # will store which agent is currently selected

def init_and_extend_messages(system_msg: object, history: list[dict[str, str]]):
    messages = [system_msg]
    messages.extend(history)
    return messages

# --- Respond function ---
def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p
):
    """Routes the user message to the active agent."""
    global active_agent

    if not active_agent:
        return "⚠️ Please select an agent before chatting."



    try:
        # Route message to correct agent
        if active_agent == "stored":
            response = init_stored_image_agent().run(
                f"""
                {message}
                """,
                images=stored_images
            )
        elif active_agent == "dynamic":
            json_response = main()
            print("JSON RESPONSE:", json_response)
            if isinstance(json_response, str):
                return json_response
            else:
                try:
                    response = json.dumps(json_response, indent=4)
                    return response
                except (json.JSONDecodeError, TypeError):
                    # JSONDecodeError for invalid JSON format in a string
                    # TypeError if the input is not a string or bytes-like object
                    print("Error parsing json response:", json_response)
                    return "Error generating response"
        elif active_agent == "multi":
            manager_agent = define_multi_agent()
            manager_agent.visualize()

            json_response = manager_agent.run(f"""{message}""")

            manager_agent.python_executor.state["fig"]
            if isinstance(json_response, str):
                response = json_response
                return response
            else:
                try:
                    response = json.dumps(json_response, indent=4)
                    return response
                except (json.JSONDecodeError, TypeError):
                    # JSONDecodeError for invalid JSON format in a string
                    # TypeError if the input is not a string or bytes-like object
                    print("Error parsing json response:", json_response)
                    return "Error generating response"
        elif active_agent == "nsfw check":
            json_response = classify_image_if_nsfw(message)
            response = json.dumps(json_response, indent=4)
            return response
        else:
            response = f"Unknown agent selected."
    except Exception as e:
        print("Exception:", str(e))

        response = f"⚠️ Error: {e}"

    return response


# --- Button Handlers ---
def use_stored_image_agent():
    global active_agent
    active_agent = "stored"
    return "✅ Switched to Stored Image Agent."


def use_dynamic_image_agent():
    global active_agent
    active_agent = "dynamic"
    return "✅ Switched to Dynamic Image Agent."


def use_multi_agent():
    global active_agent
    active_agent = "multi"
    return "✅ Switched to Multi-Agent mode."


def use_nsfw_check():
    global active_agent
    active_agent = "nsfw check"
    classify_image_if_nsfw("https://static.api4.ai/api4.ai/nsfw/demo-pic-1.jpg")
    return "✅ Switched to NSFW Check mode."


# --- Chat Interface ---
chatbot = gr.ChatInterface(
    fn=respond,
    type="messages",
    additional_inputs=[]
)

# --- Layout ---
with gr.Blocks() as demo:

    chatbot.render()
    gr.Markdown("### 🧩 Choose an Agent:")
    with gr.Row():
        stored_img_button = gr.Button("Checked Stored Superheros", variant="secondary")
        dynamic_img_button = gr.Button("Dynamically look for superheros", variant="primary")
        multi_agent_button = gr.Button("Search superhero's using multiple agents", variant="secondary")
        check_nsfw_button = gr.Button("NSFW Check on Image", variant="stop")

    # Display agent switch confirmation message
    status_box = gr.Textbox(label="Agent Status", interactive=False)

    stored_img_button.click(fn=use_stored_image_agent, inputs=None, outputs=status_box)
    dynamic_img_button.click(fn=use_dynamic_image_agent, inputs=None, outputs=status_box)
    multi_agent_button.click(fn=use_multi_agent, inputs=None, outputs=status_box)
    check_nsfw_button.click(fn= use_nsfw_check, inputs=None, outputs=status_box)

# --- Run app ---
if __name__ == "__main__":
    demo.launch()