from dotenv import load_dotenv

import gradio as gr
import traceback
import os
import html

from agent import KnowledgeMapperAgent

load_dotenv()

def process_text_with_agent(selected_api: str, api_key: str, model: str, input_text: str):
    """
    This is the final production function. It initializes and runs the KnowledgeMapperAgent,
    then returns the interactive graph, a download path, and the summary.
    """
    # 1. Validate user inputs
    if not input_text.strip():
        error_msg = "<p style='color:orange; text-align:center;'>Please enter some text to generate a map.</p>"
        return error_msg, None, "Input text is missing."
    if not api_key.strip():
        api_key = os.getenv(f"{selected_api.upper()}_API_KEY")
        if not api_key:
            error_msg = f"<p style='color:orange; text-align:center;'>Please enter your {selected_api} API key.</p>"
            return error_msg, None, "API key is missing."

    try:
        # 2. Initialize and run the agent
        print("Initializing and running the agent...")
        agent = KnowledgeMapperAgent(
            selected_api=selected_api,
            api_key=api_key,
            model=model
        )
        final_state = agent(input_text)

        # 3. Correctly extract all three outputs from the agent's final state
        html_content = final_state.get("image_path")
        file_path = final_state.get("download_path") # This gets the file path (e.g., "concept_map.html")
        summary = final_state.get("summary", "No summary was generated.")

        # 4. Prepare the outputs for Gradio
        if html_content and not html_content.strip().startswith("<p>"):
            print("Successfully retrieved data from agent.")
            escaped_html = html.escape(html_content)
            iframe_wrapper = f'<iframe srcdoc="{escaped_html}" width="100%" height="700px" frameborder="0"></iframe>'
            
            # 5. Return the tuple in the correct order: (HTML content, file path, summary)
            return iframe_wrapper, file_path, summary
        
        else: # Handle cases where the agent returns an error or no graph
             return html_content or "<p>No graph was generated.</p>", None, summary

    except Exception as e:
        print(f"An error occurred during agent processing: {e}")
        traceback.print_exc()
        error_message = f"<p style='color:red; text-align:center;'>An unexpected error occurred: {e}</p>"
        return error_message, None, "Summary could not be generated due to an error."

def main():
    # Define the Gradio Interface
    with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as text_to_map_app:
        gr.Markdown("# Text to Knowledge Map AI Agent 🧠")
        gr.Markdown(
            "Select an API provider, add the API key, choose the model, and enter your text/question. "
            "This AI agent will help you visualize its key concepts in a non-linear format."
        )

        with gr.Row():
            with gr.Column(scale=1):
                # --- Input Components ---
                api_dropdown = gr.Dropdown(
                    choices=["Gemini", "Groq", "OpenAI", "Anthropic", "Mistral", "Cohere"],
                    label="Select LLM API Provider",
                    value="Gemini",
                    interactive=True,
                )

                text_api_box = gr.Textbox(
                    lines=1,
                    label="Enter your API key here",
                    interactive=True,
                    type="password",
                    placeholder="e.g., sk-xxxxxxxxxxxxxx"
                )

                text_model_box = gr.Textbox(
                    lines=1,
                    label="Enter model name",
                    interactive=True,
                    value="gemini-2.0-flash",
                    placeholder="e.g., gemini-1.5-flash, llama3-70b-8192"
                )

                text_input_box = gr.Textbox(
                    lines=15,
                    label="Enter your text, topic, or question",
                    placeholder="Paste a document, ask a question like 'What is quantum computing?', or enter a topic like 'The history of ancient Rome.'",
                    interactive=True,
                )

                process_button = gr.Button("Generate Knowledge Map", variant="primary")

            with gr.Column(scale=2):
                with gr.Row(equal_height=False):
                    # The first column will be for the graph
                    with gr.Column(scale=3):
                        conceptual_map_html = gr.HTML(
                            label="Interactive Knowledge Map",
                            value="<p style='text-align:center;'>The interactive map will appear here.</p>",
                        )

                    # The second column will hold the summary and download link
                    with gr.Column(scale=2):
                        summary_output = gr.Markdown(
                            label="Knowledge Map Summary",
                            value="The summary of the map will appear here."
                        )
                        download_link = gr.File(
                            label="Download HTML File",
                            interactive=False
                        )


        # Connect the button to the function and all inputs/outputs
        process_button.click(
            fn=process_text_with_agent,
            inputs=[api_dropdown, text_api_box, text_model_box, text_input_box],
            outputs=[conceptual_map_html, download_link, summary_output],
        )

    text_to_map_app.launch()


if __name__ == "__main__":
    main()