import os
import json
from typing import Optional
import gradio as gr
from gradio import Interface, Blocks
import networkx as nx
import pyvis
from pyvis.network import Network
from smolagents import CodeAgent, HfApiModel, tool, GradioUI
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from openinference.instrumentation.smolagents import SmolagentsInstrumentor

# Set up telemetry
PHOENIX_API_KEY = os.getenv("PHOENIX_API_KEY")
api_key = f"api_key={PHOENIX_API_KEY}"
os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = api_key
os.environ["PHOENIX_CLIENT_HEADERS"] = api_key
os.environ["PHOENIX_COLLECTOR_ENDPOINT"] = "https://app.phoenix.arize.com"

# Updated endpoint from local to cloud
endpoint = "https://app.phoenix.arize.com/v1/traces"
trace_provider = TracerProvider()
trace_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(endpoint)))
SmolagentsInstrumentor().instrument(tracer_provider=trace_provider)

# Define examples
examples = [
    ["Highlight the Medici family's central role in the network using all three centrality metrics."],
    ["Focus on the Strozzi family's role in the network using betweenness centrality."],
    ["Compare family positions in the network using degree and closeness centrality."],
    ["Visualize the overall network structure and the relative positions of the families."],
    ["Explore the roles of the Florentine families with the highest centrality values."]
]

class GradioUIWithExamples(GradioUI):
    def __init__(self, agent, examples=None, **kwargs):
        super().__init__(agent, **kwargs)
        self.examples = examples

    def build_interface(self):
        with gr.Blocks() as demo:
            gr.Markdown("## Florentine Families Network Analysis")

            # Main Input/Output
            input_box = gr.Textbox(
                label="Your Question",
                placeholder="Type your question about the Florentine Families graph...",
            )
            output_box = gr.Textbox(
                label="Agent's Response",
                placeholder="Response will appear here...",
                interactive=False,
            )
            submit_button = gr.Button("Submit")
            
            # Link submit button to agent logic
            submit_button.click(
                self.agent.run,
                inputs=input_box,
                outputs=output_box,
            )

            # Add Examples
            if self.examples:
                gr.Markdown("### Examples")
                for example in self.examples:
                    gr.Button(example[0]).click(
                        lambda x=example[0]: x,  # Populate input box
                        inputs=[],
                        outputs=input_box,
                    )
        return demo

    def launch(self):
        # Use the custom-built interface instead of the base class's logic
        demo = self.build_interface()
        demo.launch()

# Define a tool for analyzing the Florentine Families graph
@tool
def analyze_florentine_graph(metric: str) -> str:
    """
    Analyzes the Florentine Families graph based on the chosen centrality metric.

    Args:
        metric: The centrality metric to calculate. Valid options: 'degree', 'betweenness', or 'closeness'.

    Returns:
        A textual analysis of the chosen metric for the graph with formatted numerical results.
    """
    graph = nx.florentine_families_graph()

    if metric == "degree":
        centrality = nx.degree_centrality(graph)
    elif metric == "betweenness":
        centrality = nx.betweenness_centrality(graph)
    elif metric == "closeness":
        centrality = nx.closeness_centrality(graph)
    else:
        return "Invalid metric. Please choose 'degree', 'betweenness', or 'closeness'."

    analysis = f"Analysis of {metric} centrality:\n"
    for node, value in centrality.items():
        analysis += f"- {node}: {value:.3f}\n"
    return analysis

@tool
def save_html_to_file(html_content: str, file_path: str) -> str:
    """
    Saves the provided HTML content to a file.

    Args:
        html_content: The HTML content to save.
        file_path: The path where the HTML file will be saved.

    Returns:
        A confirmation message upon successful saving.
    """
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(html_content)
    return f"HTML content successfully saved to {file_path}"

@tool
def read_html_from_file(file_path: str) -> str:
    """
    Reads HTML content from a file.

    Args:
        file_path: The path of the HTML file to read.

    Returns:
        The HTML content as a string.
    """
    with open(file_path, 'r', encoding='utf-8') as file:
        html_content = file.read()
    return html_content

@tool
def export_graph_to_json(graph_data: dict) -> str:
    """
    Exports a NetworkX graph represented as a dictionary to JSON format.

    Args:
        graph_data: The graph data in node-link format.

    Returns:
        str: The JSON representation of the graph.
    """
    graph = nx.node_link_graph(graph_data)
    json_output = json.dumps(nx.node_link_data(graph), indent=4)
    return json_output

# Initialize the agent with proper configuration
model = HfApiModel()
agent = CodeAgent(
    tools=[analyze_florentine_graph, save_html_to_file, read_html_from_file, export_graph_to_json],
    model=model,
    additional_authorized_imports=["gradio","networkx","community_louvain","pyvis","matplotlib","json"],
    add_base_tools=True
)

# Use the fixed GradioUIWithExamples
interface = GradioUIWithExamples(agent, examples=examples)
interface.launch()