"""
Manual instrumentation client for AI tracing.
"""
import json
import uuid
from uuid import UUID
from typing import Optional

import psycopg

from nexus_library.config import DB_BACKEND_URL, get_db_url


class Span:
    """
    Manual instrumentation span for AI tracing.
    Similar to LangChain callback handler but for manual tracing.
    """
    def __init__(self, name: str, event_type: str, parent_node=None, client_ref=None, input_data=None):
        self.run_id = uuid.uuid4()
        self.name = name
        self.event_type = event_type  # Category: chain, llm, tool, agent, etc.
        self.parent_node = parent_node
        self.parent_run_id = parent_node.run_id if parent_node else None
        self.client_ref = client_ref  # Reference to NexusClient to update dropback_node
        self.input = input_data
        self.output = None
        self.error = None
        self.status = "started"
        
    def __enter__(self):
        if self.client_ref is not None:
            self.client_ref.dropback_node = self  # Update NexusClient's dropback_node
        return self
        
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            self.error = str(exc_value)
            self.status = "error"
            print(f"🔴 Error in span: {self.name} - {self.error}")
        else:
            self.status = "completed"
            print(f"✅ Completed span: {self.name}")
        
        # Restore parent as dropback
        if self.client_ref is not None:
            self.client_ref.dropback_node = self.parent_node
        
        # Add event to client's events array
        if self.client_ref is not None:
            self.client_ref.events.append({
                'run_id': str(self.run_id),
                'user_id': self.client_ref.user_id,
                'trace_id': self.client_ref.trace_id,
                'parent_run_id': str(self.parent_run_id) if self.parent_run_id else None,
                'event_type': self.event_type,
                'name': self.name,
                'input': self.input,
                'output': self.output,
                'error': self.error
            })
        
        return False  # Don't suppress exceptions
    
    def set_output(self, output):
        """Set the output of the span"""
        self.output = output
    
    def set_error(self, error):
        """Set an error for the span"""
        self.error = str(error)
        self.status = "error"


class NexusClient:
    """
    Client for manual AI tracing instrumentation.
    Similar to LangChain PrintingHandler but for manual tracing.
    """
    def __init__(self, api_key: str, db_url: Optional[str] = None):
        """
        Initialize NexusClient with API key.
        
        Args:
            api_key: API key that will be used as user_id
            db_url: Optional database URL. If not provided, uses config default.
        """
        # Use api_key as user_id (convert to UUID if it's a string)
        self.user_id = UUID(api_key) if isinstance(api_key, str) else api_key
        # Generate trace_id
        self.trace_id = str(uuid.uuid4())
        self.dropback_node = None  # Current active span (for parent tracking)
        self.events = []  # Array to store events before flushing
        self.db_url = db_url or get_db_url()
    
    def _safe_json_serialize(self, obj):
        """Safely serialize objects to JSON, converting non-serializable objects to strings"""
        try:
            return json.dumps(obj, default=str)
        except (TypeError, ValueError) as e:
            # If serialization fails, convert to string representation
            return json.dumps(str(obj))
    
    def _insert_log(self, run_id, parent_run_id, event_type, name=None, input_data=None, output_data=None, error=None):
        """Helper method to insert log into database"""
        if not self.db_url:
            print(f"\033[91m🔴 Database URL not configured. Set NEXUS_DB_BACKEND_URL environment variable or pass db_url to NexusClient.\033[0m")
            return
            
        try:
            with psycopg.connect(self.db_url) as conn:
                with conn.cursor() as cur:
                    # Convert input/output to JSON strings
                    input_json = self._safe_json_serialize(input_data) if input_data is not None else None
                    output_json = self._safe_json_serialize(output_data) if output_data is not None else None
                    
                    # parent_run_id can be NULL for root/top-level invocations
                    
                    cur.execute(
                        """
                        INSERT INTO agent_logs 
                        (run_id, user_id, trace_id, parent_run_id, event_type, name, input, output, error)
                        VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s)
                        ON CONFLICT (run_id) DO UPDATE SET
                            event_type = EXCLUDED.event_type,
                            name = COALESCE(EXCLUDED.name, agent_logs.name),
                            input = COALESCE(EXCLUDED.input, agent_logs.input),
                            output = COALESCE(EXCLUDED.output, agent_logs.output),
                            error = COALESCE(EXCLUDED.error, agent_logs.error)
                        """,
                        (
                            UUID(run_id),
                            self.user_id,
                            self.trace_id,
                            UUID(parent_run_id) if parent_run_id else None,
                            event_type,
                            name,
                            input_json,
                            output_json,
                            str(error) if error else None
                        )
                    )
                    conn.commit()
        except Exception as e:
            print(f"\033[91m🔴 Database insert error: {e}\033[0m")
        
    def span(self, name: str, event_type: str, input_data=None):
        """
        Create a new span for manual instrumentation.
        
        Args:
            name: Display name (e.g., "get_weather", "gemini-2.5-flash")
            event_type: Category (chain, llm, tool, agent, retriever, etc.)
            input_data: Optional input data for the span
        """
        return Span(name, event_type, parent_node=self.dropback_node, 
                   client_ref=self, input_data=input_data)
    
    def flush(self):
        """
        Flush events to database.
        """
        num_events = len(self.events)
        
        for event in self.events:
            # Insert into database
            self._insert_log(
                run_id=event['run_id'],
                parent_run_id=event['parent_run_id'],
                event_type=event['event_type'],
                name=event['name'],
                input_data=event['input'],
                output_data=event['output'],
                error=event['error']
            )
        
        # Clear events after flushing
        self.events.clear()
        print(f"✅ Flush complete. {num_events} events pushed to database.")
    
    def get_all_events(self, trace_id: Optional[str] = None):
        """
        Retrieve all events for this client's user_id.
        
        Args:
            trace_id: Optional trace_id to filter by. If not provided, uses client's trace_id.
        
        Returns:
            List of dictionaries containing event data
        """
        if not self.db_url:
            print(f"\033[91m🔴 Database URL not configured.\033[0m")
            return []
        
        target_trace_id = trace_id or self.trace_id
        
        try:
            with psycopg.connect(self.db_url) as conn:
                with conn.cursor() as cur:
                    query = """
                        SELECT run_id, user_id, trace_id, parent_run_id, event_type, name, input, output, error
                        FROM agent_logs
                        WHERE user_id = %s AND trace_id = %s
                        ORDER BY run_id
                    """
                    cur.execute(query, (self.user_id, target_trace_id))
                    rows = cur.fetchall()
                    
                    # Convert to list of dictionaries
                    events = []
                    for row in rows:
                        run_id, uid, tid, parent_run_id, event_type, name, input_data, output_data, error = row
                        events.append({
                            'run_id': str(run_id),
                            'user_id': str(uid),
                            'trace_id': tid,
                            'parent_run_id': str(parent_run_id) if parent_run_id else None,
                            'event_type': event_type,
                            'name': name,
                            'input': input_data,
                            'output': output_data,
                            'error': error
                        })
                    
                    return events
                    
        except Exception as e:
            print(f"\033[91m🔴 Error retrieving events: {e}\033[0m")
            import traceback
            traceback.print_exc()
            return []
    
    def print_tree_structure(self, trace_id: Optional[str] = None):
        """
        Display events as a tree structure for this client's user_id.
        
        Args:
            trace_id: Optional trace_id to filter by. If not provided, uses client's trace_id.
        """
        events = self.get_all_events(trace_id=trace_id)
        
        if not events:
            print("No events found")
            return
        
        # Build parent-child mapping
        children_map = {}
        event_map = {}
        
        # First pass: build event map
        for event in events:
            run_id = event['run_id']
            event_map[run_id] = event
        
        # Second pass: build children map and identify root nodes
        root_nodes = []
        for event in events:
            run_id = event['run_id']
            parent_id = event['parent_run_id']
            
            if parent_id and parent_id in event_map:
                # Has a valid parent in our dataset
                if parent_id not in children_map:
                    children_map[parent_id] = []
                children_map[parent_id].append(run_id)
            else:
                # Root node (no parent or parent not in dataset)
                root_nodes.append(run_id)
        
        # Color codes for terminal
        colors = {
            'chain': '\033[94m',      # Blue
            'llm': '\033[92m',        # Green
            'chat_model': '\033[92m', # Green
            'tool': '\033[91m',       # Red
            'agent': '\033[93m',      # Yellow
            'retriever': '\033[96m',  # Cyan
            'text': '\033[95m',       # Magenta
            'retry': '\033[93m',      # Yellow
        }
        reset = '\033[0m'
        error_color = '\033[91m'  # Red for errors
        
        def print_node(node_id, prefix="", is_last=True, is_root=False):
            """Recursively print tree nodes"""
            if node_id not in event_map:
                return
            
            event = event_map[node_id]
            event_type = event['event_type']
            name = event['name'] or ''
            has_error = event.get('error') is not None
            
            # Choose color based on event type
            color = colors.get(event_type, reset)
            
            # Build the tree branch
            if is_root:
                branch = "└── "
            else:
                branch = "└── " if is_last else "├── "
            
            # Build label
            label_parts = [event_type]
            if name:
                label_parts.append(f"({name})")
            if has_error:
                label_parts.append("⚠️ ERROR")
            
            label = " ".join(label_parts)
            
            # Print the node
            print(f"{prefix}{branch}{color}{label}{reset}")
            
            # Get children
            children = sorted(children_map.get(node_id, []))
            
            # Print children
            for i, child_id in enumerate(children):
                is_last_child = (i == len(children) - 1)
                extension = "    " if is_last else "│   "
                print_node(child_id, prefix + extension, is_last_child, is_root=False)
        
        print(f"\n🌳 Execution Tree Structure")
        print(f"{'='*60}\n")
        
        if not root_nodes:
            print("No root nodes found")
            return
        
        # Print each root tree
        for i, root_id in enumerate(sorted(root_nodes)):
            if i > 0:
                print()  # Space between multiple root trees
            print_node(root_id, "", is_last=(i == len(root_nodes) - 1), is_root=True)
        
        print(f"\n{'='*60}")
        print(f"📊 Total events: {len(events)}")
        print(f"🌲 Root nodes: {len(root_nodes)}")
    
    def build_tree_data_from_events(self, trace_id: Optional[str] = None):
        """
        Build tree_data structure from database events for this client's user_id.
        
        Args:
            trace_id: Optional trace_id to filter by. If not provided, uses client's trace_id.
        
        Returns:
            tree_data dictionary
        """
        tree_data = {}
        all_events = self.get_all_events(trace_id=trace_id)
        
        for event in all_events:
            if event['parent_run_id'] is None:
                tree_data[event['run_id']] = {
                    'name': event['event_type'],
                    'display_name': event.get('name'),
                    'input': event.get('input'),
                    'output': event.get('output'),
                    'error': event.get('error'),
                    'children': []
                }
            else:
                # Ensure parent exists in tree_data
                if event['parent_run_id'] not in tree_data:
                    tree_data[event['parent_run_id']] = {
                        'name': 'unknown',
                        'display_name': None,
                        'input': None,
                        'output': None,
                        'error': None,
                        'children': []
                    }
                tree_data[event['parent_run_id']]['children'].append(event['run_id'])
                tree_data[event['run_id']] = {
                    'name': event['event_type'],
                    'display_name': event.get('name'),
                    'input': event.get('input'),
                    'output': event.get('output'),
                    'error': event.get('error'),
                    'children': []
                }
        
        return tree_data
    
    def _tree_data_to_agent_string(self, tree_data):
        """
        Convert tree_data structure into a well-formatted string that can be easily
        interpreted by an agent/LLM.
        
        Args:
            tree_data: Dictionary where keys are run_ids and values are dicts with
                      'name', 'display_name', 'input', 'output', 'error', 'children'
        
        Returns:
            Formatted string representation of the execution tree
        """
        def format_node(node_id, node, depth=0, is_last=True, prefix=""):
            """Recursively format a node and its children"""
            indent = "  " * depth
            name = node.get('name', 'unknown')
            display_name = node.get('display_name')
            input_data = node.get('input')
            output_data = node.get('output')
            error_data = node.get('error')
            children = node.get('children', [])
            
            # Build node header
            parts = []
            node_label = name
            if display_name:
                node_label += f" ({display_name})"
            if error_data:
                node_label += " [ERROR]"
            
            parts.append(f"{prefix}{indent}Node: {node_label}")
            parts.append(f"{prefix}{indent}Run ID: {node_id}")
            
            # Add input
            if input_data:
                input_str = json.dumps(input_data, default=str, indent=2)
                if len(input_str) > 500:
                    input_str = input_str[:500] + "... [truncated]"
                parts.append(f"{prefix}{indent}Input: {input_str}")
            
            # Add output
            if output_data:
                output_str = json.dumps(output_data, default=str, indent=2)
                if len(output_str) > 500:
                    output_str = output_str[:500] + "... [truncated]"
                parts.append(f"{prefix}{indent}Output: {output_str}")
            
            # Add error
            if error_data:
                error_str = str(error_data)
                if len(error_str) > 500:
                    error_str = error_str[:500] + "... [truncated]"
                parts.append(f"{prefix}{indent}Error: {error_str}")
            
            # Add children count
            if children:
                parts.append(f"{prefix}{indent}Children: {len(children)} child node(s)")
            
            parts.append("")  # Empty line for readability
            
            # Process children
            for i, child_id in enumerate(children):
                if child_id in tree_data:
                    child_node = tree_data[child_id]
                    is_last_child = (i == len(children) - 1)
                    child_prefix = prefix
                    parts.extend(format_node(child_id, child_node, depth + 1, is_last_child, child_prefix))
            
            return parts
        
        # Find root nodes
        all_children = set()
        for node in tree_data.values():
            all_children.update(node.get('children', []))
        
        root_nodes = [node_id for node_id in tree_data.keys() if node_id not in all_children]
        
        # Build the output string
        output_parts = []
        output_parts.append("=" * 80)
        output_parts.append("EXECUTION TREE STRUCTURE")
        output_parts.append("=" * 80)
        output_parts.append(f"Total Nodes: {len(tree_data)}")
        output_parts.append(f"Root Nodes: {len(root_nodes)}")
        output_parts.append("")
        
        # Format each root tree
        for i, root_id in enumerate(sorted(root_nodes)):
            if i > 0:
                output_parts.append("")
                output_parts.append("-" * 80)
                output_parts.append("")
            
            root_node = tree_data[root_id]
            output_parts.extend(format_node(root_id, root_node, depth=0, is_last=True, prefix=""))
        
        output_parts.append("=" * 80)
        
        return "\n".join(output_parts)
    
    def get_trace_tree_string(self, trace_id: Optional[str] = None):
        """
        Get execution tree as an agent-readable string for this client's user_id.
        
        Args:
            trace_id: Optional trace_id to filter by. If not provided, uses client's trace_id.
        
        Returns:
            Formatted string representation of the execution tree
        """
        tree_data = self.build_tree_data_from_events(trace_id=trace_id)
        return self._tree_data_to_agent_string(tree_data)
    
    def print_tree_data(self, tree_data):
        """
        Print tree_data structure as a repository-style tree.
        
        Args:
            tree_data: Dictionary where keys are run_ids and values are dicts with
                      'name' and 'children' (list of run_ids)
        """
        # Color codes for different event types
        colors = {
            'chain': '\033[94m',      # Blue
            'llm': '\033[92m',        # Green
            'chat_model': '\033[92m', # Green
            'tool': '\033[91m',       # Red
            'agent': '\033[93m',      # Yellow
            'retriever': '\033[96m',  # Cyan
            'text': '\033[95m',       # Magenta
            'retry': '\033[93m',      # Yellow
        }
        reset = '\033[0m'
        dim = '\033[90m'  # Gray for metadata
        error_color = '\033[91m'  # Red for errors
        
        def format_data(data, max_length=80):
            """Format JSON data for display, truncating if too long"""
            if data is None:
                return None
            try:
                if isinstance(data, dict):
                    data_str = json.dumps(data, default=str)
                else:
                    data_str = str(data)
                
                if len(data_str) > max_length:
                    return data_str[:max_length] + "..."
                return data_str
            except:
                return str(data)[:max_length] + "..." if len(str(data)) > max_length else str(data)
        
        def print_node(node_id, prefix="", is_last=True, is_root=False):
            """Recursively print tree nodes"""
            if node_id not in tree_data:
                return
            
            node = tree_data[node_id]
            name = node.get('name', 'unknown')
            display_name = node.get('display_name')
            input_data = node.get('input')
            output_data = node.get('output')
            error_data = node.get('error')
            children = node.get('children', [])
            
            # Choose color based on event type
            color = colors.get(name, reset)
            
            # Build the tree branch
            if is_root:
                branch = ""
                child_prefix = ""
            else:
                branch = "└── " if is_last else "├── "
                child_prefix = "    " if is_last else "│   "
            
            # Build label with display name if available
            label_parts = [name]
            if display_name:
                label_parts.append(f"({display_name})")
            if error_data:
                label_parts.append(f"{error_color}⚠️ ERROR{reset}")
            
            label = " ".join(label_parts)
            
            # Print the node
            print(f"{prefix}{branch}{color}{label}{reset}")
            
            # Print input/output/error details
            detail_prefix = prefix + child_prefix
            if input_data:
                input_str = format_data(input_data)
                print(f"{detail_prefix}{dim}  📥 Input: {input_str}{reset}")
            if output_data:
                output_str = format_data(output_data)
                print(f"{detail_prefix}{dim}  📤 Output: {output_str}{reset}")
            if error_data:
                error_str = format_data(error_data)
                print(f"{detail_prefix}{error_color}  ❌ Error: {error_str}{reset}")
            
            # Print children
            for i, child_id in enumerate(children):
                is_last_child = (i == len(children) - 1)
                extension = "    " if is_last else "│   "
                print_node(child_id, prefix + extension, is_last_child, is_root=False)
        
        # Find root nodes (nodes that are not children of any other node)
        all_children = set()
        for node in tree_data.values():
            all_children.update(node.get('children', []))
        
        root_nodes = [node_id for node_id in tree_data.keys() if node_id not in all_children]
        
        print(f"\n🌳 Tree Structure")
        print(f"{'='*60}\n")
        
        if not root_nodes:
            print("No root nodes found")
            return
        
        # Print each root tree
        for i, root_id in enumerate(sorted(root_nodes)):
            if i > 0:
                print()  # Space between multiple root trees
            print_node(root_id, "", is_last=(i == len(root_nodes) - 1), is_root=True)
        
        print(f"\n{'='*60}")
        print(f"📊 Total nodes: {len(tree_data)}")
        print(f"🌲 Root nodes: {len(root_nodes)}")