"""
Tool Graph Web UI

A Flask-based web application for visualizing and exploring the tool composition graph
generated by ToolGraphComposer.
"""

import json
import os
from flask import Flask, render_template, jsonify, request
import pickle
from typing import Dict, Any, List, Optional


class ToolGraphWebUI:
    """Web interface for visualizing tool composition graphs."""

    def __init__(self, graph_data_path: Optional[str] = None):
        self.app = Flask(__name__, template_folder="templates", static_folder="static")
        self.graph_data: Optional[Dict[str, Any]] = None
        self.graph_data_path = graph_data_path

        self._setup_routes()
        self._load_graph_data()

    def _setup_routes(self):
        """Setup Flask routes."""

        @self.app.route("/")
        def index():
            """Main graph visualization page."""
            return render_template("tool_graph.html")

        @self.app.route("/api/graph")
        def get_graph():
            """API endpoint to get graph data."""
            if not self.graph_data:
                return jsonify({"error": "No graph data loaded"}), 404

            # Convert to format suitable for D3.js
            d3_data = self._convert_to_d3_format(self.graph_data)
            return jsonify(d3_data)

        @self.app.route("/api/stats")
        def get_stats():
            """API endpoint to get graph statistics."""
            if not self.graph_data:
                return jsonify({"error": "No graph data loaded"}), 404

            return jsonify(self.graph_data.get("metadata", {}))

        @self.app.route("/api/tool/<tool_id>")
        def get_tool_details(tool_id):
            """API endpoint to get detailed information about a specific tool."""
            if not self.graph_data:
                return jsonify({"error": "No graph data loaded"}), 404

            # Find the tool
            tool = None
            for node in self.graph_data.get("nodes", []):
                if node.get("id") == tool_id:
                    tool = node
                    break

            if not tool:
                return jsonify({"error": "Tool not found"}), 404

            # Get connections
            connections = self._get_tool_connections(tool_id)

            return jsonify({"tool": tool, "connections": connections})

        @self.app.route("/api/search")
        def search_tools():
            """API endpoint to search tools."""
            query = request.args.get("q", "").lower()
            category = request.args.get("category", "")

            if not self.graph_data:
                return jsonify({"error": "No graph data loaded"}), 404

            results = []
            for node in self.graph_data.get("nodes", []):
                # Search in name and description
                if (
                    query in node.get("name", "").lower()
                    or query in node.get("description", "").lower()
                ):
                    if not category or node.get("category") == category:
                        results.append(
                            {
                                "id": node.get("id"),
                                "name": node.get("name"),
                                "description": node.get("description"),
                                "category": node.get("category"),
                            }
                        )

            return jsonify(results)

        @self.app.route("/api/load_graph", methods=["POST"])
        def load_graph():
            """API endpoint to load a new graph file."""
            data = request.json
            graph_path = data.get("path")

            if not graph_path or not os.path.exists(graph_path):
                return jsonify({"error": "Invalid graph path"}), 400

            try:
                self.graph_data_path = graph_path
                self._load_graph_data()
                return jsonify(
                    {"status": "success", "message": "Graph loaded successfully"}
                )
            except Exception as e:
                return jsonify({"error": f"Failed to load graph: {str(e)}"}), 500

    def _load_graph_data(self):
        """Load graph data from file."""
        if not self.graph_data_path:
            print("No graph data path specified")
            return

        try:
            # Try to load JSON first
            if self.graph_data_path.endswith(".json"):
                with open(self.graph_data_path, "r", encoding="utf-8") as f:
                    self.graph_data = json.load(f)
            # Try pickle format
            elif self.graph_data_path.endswith(".pkl"):
                with open(self.graph_data_path, "rb") as f:
                    self.graph_data = pickle.load(f)
                    # Handle cached format
                    if (
                        isinstance(self.graph_data, dict)
                        and "graph_data" in self.graph_data
                    ):
                        self.graph_data = self.graph_data["graph_data"]
            else:
                print(f"Unsupported file format: {self.graph_data_path}")
                return

            print(
                f"Loaded graph with {len(self.graph_data.get('nodes', []))} nodes and {len(self.graph_data.get('edges', []))} edges"
            )

        except Exception as e:
            print(f"Error loading graph data: {e}")
            self.graph_data = None

    def _convert_to_d3_format(self, graph_data: Dict[str, Any]) -> Dict[str, Any]:
        """Convert graph data to D3.js compatible format."""

        nodes = []
        links = []

        # Process nodes
        for node in graph_data.get("nodes", []):
            d3_node = {
                "id": node.get("id"),
                "name": node.get("name"),
                "description": (
                    node.get("description")[:100] + "..."
                    if len(node.get("description", "")) > 100
                    else node.get("description", "")
                ),
                "category": node.get("category"),
                "type": node.get("type"),
                "group": self._get_category_group(node.get("category")),
            }
            nodes.append(d3_node)

        # Process edges
        for edge in graph_data.get("edges", []):
            d3_link = {
                "source": edge.get("source"),
                "target": edge.get("target"),
                "value": edge.get("compatibility_score", 50),
                "confidence": edge.get("confidence", 50),
                "automation_ready": edge.get("automation_ready", False),
                "needs_transformation": edge.get("needs_transformation", False),
            }
            links.append(d3_link)

        return {
            "nodes": nodes,
            "links": links,
            "metadata": graph_data.get("metadata", {}),
        }

    def _get_category_group(self, category: str) -> int:
        """Map category to a numeric group for visualization."""
        category_groups = {
            "opentarget": 1,
            "ChEMBL": 2,
            "fda_drug_label": 3,
            "clinical_trials": 4,
            "EuropePMC": 5,
            "semantic_scholar": 6,
            "pubtator": 7,
            "monarch": 8,
            "agents": 9,
            "dataset": 10,
            "special_tools": 11,
        }
        return category_groups.get(category, 0)

    def _get_tool_connections(self, tool_id: str) -> Dict[str, List]:
        """Get incoming and outgoing connections for a tool."""
        # Guard against missing graph data
        if not self.graph_data:
            return {"incoming": [], "outgoing": []}

        incoming: List[Dict[str, Any]] = []
        outgoing: List[Dict[str, Any]] = []

        for edge in self.graph_data.get("edges", []):
            if edge.get("target") == tool_id:
                incoming.append(
                    {
                        "source": edge.get("source"),
                        "compatibility_score": edge.get("compatibility_score"),
                        "automation_ready": edge.get("automation_ready"),
                        "needs_transformation": edge.get("needs_transformation"),
                    }
                )
            elif edge.get("source") == tool_id:
                outgoing.append(
                    {
                        "target": edge.get("target"),
                        "compatibility_score": edge.get("compatibility_score"),
                        "automation_ready": edge.get("automation_ready"),
                        "needs_transformation": edge.get("needs_transformation"),
                    }
                )

        return {"incoming": incoming, "outgoing": outgoing}

    def run(self, host: str = "0.0.0.0", port: int = 5000, debug: bool = True):
        """Run the web application."""
        print(f"Starting Tool Graph Web UI on http://{host}:{port}")
        self.app.run(host=host, port=port, debug=debug)


def create_web_ui_files():
    """Create the necessary HTML, CSS, and JS files for the web UI."""

    # Create directories
    os.makedirs("templates", exist_ok=True)
    os.makedirs("static", exist_ok=True)

    # HTML template
    html_content = """<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>ToolUniverse Composition Graph</title>
    <link rel="stylesheet" href="{{ url_for('static', filename='style.css') }}">
    <script src="https://d3js.org/d3.v7.min.js"></script>
</head>
<body>
    <div id="app">
        <header>
            <h1>ToolUniverse Composition Graph</h1>
            <div class="controls">
                <div class="search-container">
                    <input type="text" id="search" placeholder="Search tools...">
                    <select id="categoryFilter">
                        <option value="">All Categories</option>
                    </select>
                </div>
                <div class="view-controls">
                    <button id="resetZoom">Reset Zoom</button>
                    <button id="showStats">Statistics</button>
                    <label>
                        <input type="checkbox" id="showEdgeLabels"> Show Edge Labels
                    </label>
                </div>
            </div>
        </header>

        <main>
            <div id="graph-container">
                <svg id="graph"></svg>
            </div>

            <div id="sidebar">
                <div id="tool-details" class="panel">
                    <h3>Tool Details</h3>
                    <div id="tool-info">
                        <p>Select a tool to view details</p>
                    </div>
                </div>

                <div id="graph-stats" class="panel" style="display: none;">
                    <h3>Graph Statistics</h3>
                    <div id="stats-content"></div>
                </div>
            </div>
        </main>
    </div>

    <script src="{{ url_for('static', filename='graph.js') }}"></script>
</body>
</html>"""

    # CSS styles
    css_content = """
body {
    margin: 0;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
    background-color: #f5f5f5;
}

#app {
    display: flex;
    flex-direction: column;
    height: 100vh;
}

header {
    background: white;
    padding: 1rem;
    border-bottom: 1px solid #ddd;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

header h1 {
    margin: 0 0 1rem 0;
    color: #333;
}

.controls {
    display: flex;
    justify-content: space-between;
    align-items: center;
    flex-wrap: wrap;
    gap: 1rem;
}

.search-container {
    display: flex;
    gap: 0.5rem;
}

.search-container input, .search-container select {
    padding: 0.5rem;
    border: 1px solid #ddd;
    border-radius: 4px;
}

.view-controls {
    display: flex;
    gap: 1rem;
    align-items: center;
}

.view-controls button {
    padding: 0.5rem 1rem;
    background: #007bff;
    color: white;
    border: none;
    border-radius: 4px;
    cursor: pointer;
}

.view-controls button:hover {
    background: #0056b3;
}

main {
    display: flex;
    flex: 1;
    overflow: hidden;
}

#graph-container {
    flex: 1;
    position: relative;
}

#graph {
    width: 100%;
    height: 100%;
    background: white;
}

#sidebar {
    width: 300px;
    background: white;
    border-left: 1px solid #ddd;
    overflow-y: auto;
    padding: 1rem;
}

.panel {
    margin-bottom: 2rem;
}

.panel h3 {
    margin-top: 0;
    color: #333;
    border-bottom: 1px solid #eee;
    padding-bottom: 0.5rem;
}

/* Graph styling */
.node {
    stroke: #fff;
    stroke-width: 2px;
    cursor: pointer;
}

.node:hover {
    stroke-width: 3px;
    stroke: #333;
}

.node.selected {
    stroke: #ff6b35;
    stroke-width: 3px;
}

.link {
    stroke: #999;
    stroke-opacity: 0.6;
    fill: none;
}

.link.high-compatibility {
    stroke: #2ecc71;
    stroke-width: 2px;
}

.link.medium-compatibility {
    stroke: #f39c12;
    stroke-width: 1.5px;
}

.link.low-compatibility {
    stroke: #e74c3c;
    stroke-width: 1px;
}

.node-label {
    font-size: 10px;
    font-family: sans-serif;
    pointer-events: none;
}

.edge-label {
    font-size: 8px;
    font-family: sans-serif;
    fill: #666;
    pointer-events: none;
}

/* Tool details styling */
.tool-info-item {
    margin-bottom: 1rem;
    padding: 0.5rem;
    background: #f8f9fa;
    border-radius: 4px;
}

.tool-info-item strong {
    display: block;
    margin-bottom: 0.25rem;
    color: #495057;
}

.connections {
    margin-top: 1rem;
}

.connection-list {
    max-height: 200px;
    overflow-y: auto;
}

.connection-item {
    padding: 0.25rem;
    border-bottom: 1px solid #eee;
    font-size: 0.9em;
}

.compatibility-score {
    float: right;
    font-weight: bold;
}

.score-high { color: #2ecc71; }
.score-medium { color: #f39c12; }
.score-low { color: #e74c3c; }

/* Responsive design */
@media (max-width: 768px) {
    main {
        flex-direction: column;
    }

    #sidebar {
        width: 100%;
        height: 40%;
        border-left: none;
        border-top: 1px solid #ddd;
    }

    .controls {
        flex-direction: column;
        align-items: stretch;
    }

    .view-controls {
        justify-content: center;
    }
}
"""

    # JavaScript for D3.js visualization
    js_content = """
class ToolGraphVisualizer {
    constructor() {
        this.svg = d3.select("#graph");
        this.width = window.innerWidth - 300; // Account for sidebar
        this.height = window.innerHeight - 100; // Account for header
        this.simulation = null;
        this.nodes = [];
        this.links = [];
        this.selectedNode = null;

        this.initializeGraph();
        this.loadGraphData();
        this.setupEventListeners();

        // Resize handler
        window.addEventListener('resize', () => this.handleResize());
    }

    initializeGraph() {
        this.svg
            .attr("width", this.width)
            .attr("height", this.height);

        // Add zoom behavior
        const zoom = d3.zoom()
            .scaleExtent([0.1, 10])
            .on("zoom", (event) => {
                this.svg.select("g").attr("transform", event.transform);
            });

        this.svg.call(zoom);

        // Create main group for zoomable content
        this.graphGroup = this.svg.append("g");

        // Create arrow markers for directed edges
        this.svg.append("defs").append("marker")
            .attr("id", "arrow")
            .attr("viewBox", "0 -5 10 10")
            .attr("refX", 20)
            .attr("refY", 0)
            .attr("markerWidth", 6)
            .attr("markerHeight", 6)
            .attr("orient", "auto")
            .append("path")
            .attr("d", "M0,-5L10,0L0,5")
            .attr("fill", "#666");
    }

    async loadGraphData() {
        try {
            const response = await fetch('/api/graph');
            const data = await response.json();

            if (data.error) {
                this.showError(data.error);
                return;
            }

            this.nodes = data.nodes;
            this.links = data.links;

            this.updateCategoryFilter();
            this.renderGraph();

        } catch (error) {
            this.showError(`Failed to load graph data: ${error.message}`);
        }
    }

    updateCategoryFilter() {
        const categories = [...new Set(this.nodes.map(n => n.category))].sort();
        const select = d3.select("#categoryFilter");

        select.selectAll("option:not(:first-child)").remove();

        categories.forEach(category => {
            select.append("option")
                .attr("value", category)
                .text(category);
        });
    }

    renderGraph() {
        // Clear existing graph
        this.graphGroup.selectAll("*").remove();

        // Create simulation
        this.simulation = d3.forceSimulation(this.nodes)
            .force("link", d3.forceLink(this.links).id(d => d.id).distance(80))
            .force("charge", d3.forceManyBody().strength(-300))
            .force("center", d3.forceCenter(this.width / 2, this.height / 2))
            .force("collision", d3.forceCollide().radius(20));

        // Create links
        const link = this.graphGroup.append("g")
            .attr("class", "links")
            .selectAll("line")
            .data(this.links)
            .enter().append("line")
            .attr("class", d => `link ${this.getCompatibilityClass(d.value)}`)
            .attr("marker-end", "url(#arrow)");

        // Create nodes
        const node = this.graphGroup.append("g")
            .attr("class", "nodes")
            .selectAll("circle")
            .data(this.nodes)
            .enter().append("circle")
            .attr("class", "node")
            .attr("r", 8)
            .attr("fill", d => this.getNodeColor(d.group))
            .call(d3.drag()
                .on("start", (event, d) => this.dragStarted(event, d))
                .on("drag", (event, d) => this.dragged(event, d))
                .on("end", (event, d) => this.dragEnded(event, d)))
            .on("click", (event, d) => this.selectNode(d));

        // Add node labels
        const label = this.graphGroup.append("g")
            .attr("class", "labels")
            .selectAll("text")
            .data(this.nodes)
            .enter().append("text")
            .attr("class", "node-label")
            .attr("dx", 12)
            .attr("dy", 4)
            .text(d => d.name);

        // Add simulation tick handler
        this.simulation.on("tick", () => {
            link
                .attr("x1", d => d.source.x)
                .attr("y1", d => d.source.y)
                .attr("x2", d => d.target.x)
                .attr("y2", d => d.target.y);

            node
                .attr("cx", d => d.x)
                .attr("cy", d => d.y);

            label
                .attr("x", d => d.x)
                .attr("y", d => d.y);
        });

        // Add tooltips
        node.append("title")
            .text(d => `${d.name}\\n${d.description}\\nCategory: ${d.category}`);
    }

    getCompatibilityClass(score) {
        if (score >= 80) return "high-compatibility";
        if (score >= 60) return "medium-compatibility";
        return "low-compatibility";
    }

    getNodeColor(group) {
        const colors = d3.schemeCategory10;
        return colors[group % colors.length];
    }

    selectNode(node) {
        // Update visual selection
        this.graphGroup.selectAll(".node")
            .classed("selected", false);

        this.graphGroup.selectAll(".node")
            .filter(d => d.id === node.id)
            .classed("selected", true);

        this.selectedNode = node;
        this.loadToolDetails(node.id);
    }

    async loadToolDetails(toolId) {
        try {
            const response = await fetch(`/api/tool/${toolId}`);
            const data = await response.json();

            if (data.error) {
                this.showError(data.error);
                return;
            }

            this.displayToolDetails(data.tool, data.connections);

        } catch (error) {
            this.showError(`Failed to load tool details: ${error.message}`);
        }
    }

    displayToolDetails(tool, connections) {
        const container = d3.select("#tool-info");
        container.html("");

        // Basic tool information
        container.append("div")
            .attr("class", "tool-info-item")
            .html(`<strong>Name:</strong> ${tool.name}`);

        container.append("div")
            .attr("class", "tool-info-item")
            .html(`<strong>Type:</strong> ${tool.type}`);

        container.append("div")
            .attr("class", "tool-info-item")
            .html(`<strong>Category:</strong> ${tool.category}`);

        container.append("div")
            .attr("class", "tool-info-item")
            .html(`<strong>Description:</strong> ${tool.description}`);

        // Connections
        const connectionsDiv = container.append("div")
            .attr("class", "connections");

        // Incoming connections
        if (connections.incoming.length > 0) {
            connectionsDiv.append("h4").text("Can receive input from:");
            const incomingList = connectionsDiv.append("div")
                .attr("class", "connection-list");

            connections.incoming.forEach(conn => {
                const item = incomingList.append("div")
                    .attr("class", "connection-item");

                item.append("span").text(conn.source);
                item.append("span")
                    .attr("class", `compatibility-score ${this.getScoreClass(conn.compatibility_score)}`)
                    .text(`${conn.compatibility_score}%`);
            });
        }

        // Outgoing connections
        if (connections.outgoing.length > 0) {
            connectionsDiv.append("h4").text("Can provide input to:");
            const outgoingList = connectionsDiv.append("div")
                .attr("class", "connection-list");

            connections.outgoing.forEach(conn => {
                const item = outgoingList.append("div")
                    .attr("class", "connection-item");

                item.append("span").text(conn.target);
                item.append("span")
                    .attr("class", `compatibility-score ${this.getScoreClass(conn.compatibility_score)}`)
                    .text(`${conn.compatibility_score}%`);
            });
        }
    }

    getScoreClass(score) {
        if (score >= 80) return "score-high";
        if (score >= 60) return "score-medium";
        return "score-low";
    }

    setupEventListeners() {
        // Search functionality
        d3.select("#search").on("input", (event) => {
            this.searchTools(event.target.value);
        });

        // Category filter
        d3.select("#categoryFilter").on("change", (event) => {
            this.filterByCategory(event.target.value);
        });

        // Reset zoom
        d3.select("#resetZoom").on("click", () => {
            this.svg.transition().duration(750).call(
                d3.zoom().transform,
                d3.zoomIdentity
            );
        });

        // Show statistics
        d3.select("#showStats").on("click", () => {
            this.toggleStats();
        });

        // Edge labels toggle
        d3.select("#showEdgeLabels").on("change", (event) => {
            this.toggleEdgeLabels(event.target.checked);
        });
    }

    searchTools(query) {
        if (!query) {
            this.graphGroup.selectAll(".node").style("opacity", 1);
            this.graphGroup.selectAll(".node-label").style("opacity", 1);
            return;
        }

        const queryLower = query.toLowerCase();
        this.graphGroup.selectAll(".node")
            .style("opacity", d =>
                d.name.toLowerCase().includes(queryLower) ||
                d.description.toLowerCase().includes(queryLower) ? 1 : 0.2
            );

        this.graphGroup.selectAll(".node-label")
            .style("opacity", d =>
                d.name.toLowerCase().includes(queryLower) ||
                d.description.toLowerCase().includes(queryLower) ? 1 : 0.2
            );
    }

    filterByCategory(category) {
        if (!category) {
            this.graphGroup.selectAll(".node").style("display", "block");
            this.graphGroup.selectAll(".node-label").style("display", "block");
            return;
        }

        this.graphGroup.selectAll(".node")
            .style("display", d => d.category === category ? "block" : "none");

        this.graphGroup.selectAll(".node-label")
            .style("display", d => d.category === category ? "block" : "none");
    }

    async toggleStats() {
        const statsPanel = d3.select("#graph-stats");
        const isVisible = statsPanel.style("display") !== "none";

        if (isVisible) {
            statsPanel.style("display", "none");
            return;
        }

        try {
            const response = await fetch('/api/stats');
            const stats = await response.json();

            this.displayStats(stats);
            statsPanel.style("display", "block");

        } catch (error) {
            this.showError(`Failed to load statistics: ${error.message}`);
        }
    }

    displayStats(stats) {
        const container = d3.select("#stats-content");
        container.html("");

        container.append("p").html(`<strong>Total Tools:</strong> ${stats.total_tools || 'N/A'}`);
        container.append("p").html(`<strong>Analysis Depth:</strong> ${stats.analysis_depth || 'N/A'}`);
        container.append("p").html(`<strong>Min Compatibility Score:</strong> ${stats.min_compatibility_score || 'N/A'}`);
        container.append("p").html(`<strong>Created:</strong> ${stats.creation_time || 'N/A'}`);
    }

    toggleEdgeLabels(show) {
        // Implementation for edge labels toggle
        // This would require additional data binding for edge labels
        console.log("Edge labels toggle:", show);
    }

    dragStarted(event, d) {
        if (!event.active) this.simulation.alphaTarget(0.3).restart();
        d.fx = d.x;
        d.fy = d.y;
    }

    dragged(event, d) {
        d.fx = event.x;
        d.fy = event.y;
    }

    dragEnded(event, d) {
        if (!event.active) this.simulation.alphaTarget(0);
        d.fx = null;
        d.fy = null;
    }

    handleResize() {
        this.width = window.innerWidth - 300;
        this.height = window.innerHeight - 100;

        this.svg
            .attr("width", this.width)
            .attr("height", this.height);

        if (this.simulation) {
            this.simulation
                .force("center", d3.forceCenter(this.width / 2, this.height / 2))
                .restart();
        }
    }

    showError(message) {
        console.error(message);
        d3.select("#tool-info").html(`<div style="color: red;">Error: ${message}</div>`);
    }
}

// Initialize the visualizer when the page loads
document.addEventListener('DOMContentLoaded', () => {
    new ToolGraphVisualizer();
});
"""

    # Write files
    with open("templates/tool_graph.html", "w", encoding="utf-8") as f:
        f.write(html_content)

    with open("static/style.css", "w", encoding="utf-8") as f:
        f.write(css_content)

    with open("static/graph.js", "w", encoding="utf-8") as f:
        f.write(js_content)

    print("Web UI files created successfully!")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        graph_path = sys.argv[1]
    else:
        graph_path = "./tool_composition_graph.json"

    # Create web UI files if they don't exist
    if not os.path.exists("templates/tool_graph.html"):
        create_web_ui_files()

    # Start the web UI
    ui = ToolGraphWebUI(graph_path)
    ui.run()
