#!/usr/bin/env python3
"""
UI Builder Recursion Visualizer - Visualization component for the recursive UI builder agent.
"""

import json
import threading
import webbrowser
import logging
from http.server import HTTPServer, BaseHTTPRequestHandler

from .ui_builder_robust_agent import UITransitionGraph, UINode, UIEdge

class UIGraphVisualizerServer:
    """Simple HTTP server to visualize the graph in real-time."""
    
    def __init__(self, initial_graph: UITransitionGraph = None):
        self.graph = initial_graph or UITransitionGraph(nodes={}, edges=[])
        self.lock = threading.Lock()
        self.httpd = None
        self.server_thread = None
        self.current_traversing_node = None  # Track the currently traversing node
        self.resource_url = None  # Resource URL for remote sandbox
        self.sandbox_loaded = False  # Track if sandbox has been loaded
        self.agent = None  # Reference to the agent for calling navigation methods
        self.start_build_requested = False  # Track if build has been requested
        self.max_depth = 1  # Default max depth
    
    def set_agent(self, agent):
        """Set the agent reference for calling navigation methods."""
        self.agent = agent
    
    def update_graph(self, graph: UITransitionGraph):
        """Update the graph data."""
        with self.lock:
            # Check if the graph object has current_traversing_node info
            if hasattr(graph, 'graph_dict'):
                # This is a TempGraph object
                graph_dict = graph.graph_dict
                if 'current_traversing_node' in graph_dict:
                    self.current_traversing_node = graph_dict['current_traversing_node']
                    # Remove the special attribute to avoid serialization issues
                    temp_graph_dict = graph_dict.copy()
                    del temp_graph_dict['current_traversing_node']
                    # Create a new UITransitionGraph from the modified dict
                    nodes = {}
                    for node_id, node_data in temp_graph_dict['nodes'].items():
                        nodes[node_id] = UINode(**node_data)
                    edges = [UIEdge(**edge_data) for edge_data in temp_graph_dict['edges']]
                    self.graph = UITransitionGraph(nodes=nodes, edges=edges)
                elif 'error_state' in graph_dict:
                    # Handle error state
                    temp_graph_dict = graph_dict.copy()
                    # Create a new UITransitionGraph from the modified dict
                    nodes = {}
                    for node_id, node_data in temp_graph_dict['nodes'].items():
                        nodes[node_id] = UINode(**node_data)
                    edges = [UIEdge(**edge_data) for edge_data in temp_graph_dict['edges']]
                    self.graph = UITransitionGraph(nodes=nodes, edges=edges)
                    # Clear current traversing node
                    self.current_traversing_node = None
                else:
                    # Create a new UITransitionGraph from the dict
                    nodes = {}
                    for node_id, node_data in graph_dict['nodes'].items():
                        nodes[node_id] = UINode(**node_data)
                    edges = [UIEdge(**edge_data) for edge_data in graph_dict['edges']]
                    self.graph = UITransitionGraph(nodes=nodes, edges=edges)
                    # Clear current traversing node
                    self.current_traversing_node = None
            else:
                # This is a regular UITransitionGraph object
                self.graph = graph
                # Clear current traversing node unless it was already set
                self.current_traversing_node = None
    
    def serve(self, host: str = "localhost", port: int = 8080):
        """Start the HTTP server."""
        import urllib.parse
        from http.server import BaseHTTPRequestHandler
        import os
        
        # Get the directory where this file is located
        current_dir = os.path.dirname(os.path.abspath(__file__))
        pages_dir = os.path.join(current_dir, "pages")
        
        visualizer = self
        
        class RequestHandler(BaseHTTPRequestHandler):
            def log_message(self, format, *args):
                # Override to disable logging HTTP requests
                pass
                
            def do_GET(self):
                parsed_path = urllib.parse.urlparse(self.path)
                path = parsed_path.path
                
                if path == "/":
                    self.serve_main_page()
                elif path == "/data":
                    self.serve_graph_data()
                elif path == "/navigate_back":
                    self.handle_navigate_back()
                elif path == "/go_home":
                    self.handle_go_home()
                elif path == "/style.css":
                    self.serve_css()
                elif path == "/script.js":
                    self.serve_js()
                else:
                    self.serve_404()
            
            def do_POST(self):
                parsed_path = urllib.parse.urlparse(self.path)
                path = parsed_path.path
                
                if path == "/navigate_back":
                    self.handle_navigate_back()
                elif path == "/go_home":
                    self.handle_go_home()
                elif path == "/start_build":
                    self.handle_start_build()
                elif path == "/update_config":
                    self.handle_update_config()
                else:
                    self.send_response(404)
                    self.end_headers()
            
            def handle_navigate_back(self):
                """Handle navigate back request."""
                if visualizer.agent:
                    try:
                        # Get current traversing node if available
                        current_node_id = visualizer.current_traversing_node
                        # Call the agent's navigate back method with current node id
                        visualizer.agent._navigate_back(current_node_id=current_node_id)
                        self.send_response(200)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps({"success": True, "message": "Navigate back executed"}).encode('utf-8'))
                    except Exception as e:
                        self.send_response(500)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps({"success": False, "error": str(e)}).encode('utf-8'))
                else:
                    self.send_response(400)
                    self.send_header('Content-type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({"success": False, "error": "Agent not set"}).encode('utf-8'))
            
            def handle_go_home(self):
                """Handle go home request."""
                if visualizer.agent:
                    try:
                        # Get current traversing node if available
                        current_node_id = visualizer.current_traversing_node
                        # Call the agent's go to home screen method with current node id
                        visualizer.agent._go_to_home_screen(current_node_id=current_node_id)
                        self.send_response(200)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps({"success": True, "message": "Go home executed"}).encode('utf-8'))
                    except Exception as e:
                        self.send_response(500)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps({"success": False, "error": str(e)}).encode('utf-8'))
                else:
                    self.send_response(400)
                    self.send_header('Content-type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({"success": False, "error": "Agent not set"}).encode('utf-8'))
            
            def handle_start_build(self):
                """Handle start build request."""
                if visualizer.agent and visualizer.resource_url:
                    try:
                        # Set flag to indicate build has been requested
                        visualizer.start_build_requested = True
                        self.send_response(200)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps({"success": True, "message": "Build started"}).encode('utf-8'))
                    except Exception as e:
                        self.send_response(500)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps({"success": False, "error": str(e)}).encode('utf-8'))
                else:
                    self.send_response(400)
                    self.send_header('Content-type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({"success": False, "error": "Agent not set or resource URL not available"}).encode('utf-8'))
            
            def handle_update_config(self):
                """Handle update configuration request."""
                try:
                    # Get the POST data
                    content_length = int(self.headers['Content-Length'])
                    post_data = self.rfile.read(content_length)
                    config_data = json.loads(post_data.decode('utf-8'))
                    
                    # Update configuration
                    if 'max_depth' in config_data:
                        visualizer.max_depth = int(config_data['max_depth'])
                    
                    self.send_response(200)
                    self.send_header('Content-type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({"success": True, "message": "Configuration updated"}).encode('utf-8'))
                except Exception as e:
                    self.send_response(500)
                    self.send_header('Content-type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({"success": False, "error": str(e)}).encode('utf-8'))
            
            def serve_main_page(self):
                try:
                    with open(os.path.join(pages_dir, "index.html"), "r", encoding="utf-8") as f:
                        html_content = f.read()
                    self.send_response(200)
                    self.send_header('Content-type', 'text/html')
                    self.end_headers()
                    self.wfile.write(html_content.encode('utf-8'))
                except FileNotFoundError:
                    self.serve_404()
            
            def serve_graph_data(self):
                with visualizer.lock:
                    graph_dict = visualizer.graph.to_dict()
                    # Add current traversing node info if available
                    if visualizer.current_traversing_node:
                        graph_dict['current_traversing_node'] = visualizer.current_traversing_node
                    
                    # Add resource URL if available
                    if visualizer.resource_url:
                        graph_dict['resource_url'] = visualizer.resource_url
                    
                    # Add error state if present
                    if hasattr(visualizer, 'error_state') and visualizer.error_state:
                        graph_dict['error_state'] = True
                        if hasattr(visualizer, 'error_message') and visualizer.error_message:
                            graph_dict['error_message'] = visualizer.error_message
                
                self.send_response(200)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                self.wfile.write(json.dumps(graph_dict).encode('utf-8'))
            
            def serve_css(self):
                try:
                    with open(os.path.join(pages_dir, "style.css"), "r", encoding="utf-8") as f:
                        css_content = f.read()
                    self.send_response(200)
                    self.send_header('Content-type', 'text/css')
                    self.end_headers()
                    self.wfile.write(css_content.encode('utf-8'))
                except FileNotFoundError:
                    self.serve_404()
            
            def serve_js(self):
                try:
                    with open(os.path.join(pages_dir, "script.js"), "r", encoding="utf-8") as f:
                        js_content = f.read()
                    self.send_response(200)
                    self.send_header('Content-type', 'application/javascript')
                    self.end_headers()
                    self.wfile.write(js_content.encode('utf-8'))
                except FileNotFoundError:
                    self.serve_404()
            
            def serve_404(self):
                self.send_response(404)
                self.send_header('Content-type', 'text/html')
                self.end_headers()
                self.wfile.write(b'<h1>404 - Page Not Found</h1>')
        
        # Disable HTTP request logging
        http_server_logger = logging.getLogger('http.server')
        http_server_logger.setLevel(logging.WARNING)
        
        # Start server in a separate thread
        def start_server():
            self.httpd = HTTPServer((host, port), RequestHandler)
            print(f"[SERVER] Starting server at http://{host}:{port}")
            self.httpd.serve_forever()
        
        self.server_thread = threading.Thread(target=start_server, daemon=True)
        self.server_thread.start()
        
        # Open browser
        webbrowser.open(f"http://{host}:{port}")
    
    def stop(self):
        """Stop the server."""
        if self.httpd:
            self.httpd.shutdown()
        if self.server_thread:
            self.server_thread.join()