"""
LangChain callback handler for automatic AI tracing.
"""
import json
import uuid
from uuid import UUID
from typing import Optional

from langchain_core.callbacks.base import BaseCallbackHandler
import psycopg

from nexus_library.config import get_db_url
from nexus_library.nexus_core import NexusClient


class PrintingHandler(BaseCallbackHandler):
    """
    LangChain callback handler for tracing AI execution.
    
    Color codes for different event types:
    - Green = LLM events (important)
    - Red = Errors and tool starts (critical)
    - Blue = Chain events (important)
    - Yellow = Agent events (important)
    - Cyan = Retriever events (less important)
    - Magenta = Text events (less important)
    - White = Tool end (less important)
    """
    
    def __init__(self, active_client: NexusClient, session_id: Optional[str] = None, db_url: Optional[str] = None):
        """
        Initialize PrintingHandler 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.
        """
        super().__init__()
        # Use api_key as user_id (convert to UUID if it's a string)
        self.user_id = active_client.user_id
        self.session_id = session_id
        # Generate trace_id
        self.trace_id = str(uuid.uuid4())
        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 _get_event_category(self, event_type):
        """Extract the category from event_type (e.g., 'llm_start' -> 'llm', 'tool_end' -> 'tool')"""
        # Map of event type patterns to categories
        categories = ['chain', 'llm', 'chat_model', 'retriever', 'tool', 'agent', 'retry', 'text']
        
        # Check for each category
        event_lower = event_type.lower()
        for category in categories:
            if category in event_lower:
                return category
        
        # Default fallback
        return event_type.split('_')[0] if '_' in event_type else event_type
    
    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 PrintingHandler.\033[0m")
            return
            
        try:
            with psycopg.connect(self.db_url) as conn:
                with conn.cursor() as cur:
                    # Extract category from event_type (e.g., 'llm_start' -> 'llm')
                    event_category = self._get_event_category(event_type)
                    
                    # 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 = agent_logs.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)
                        """,
                        (
                            run_id,
                            self.user_id,
                            self.trace_id,
                            parent_run_id,
                            event_category,  # Use category instead of full 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 on_llm_start(self, serialized, prompts, run_id, parent_run_id=None, tags=None, metadata=None, **kwargs):
        # Extract model name from serialized
        model_name = None
        if isinstance(serialized, dict):
            model_name = serialized.get('name') or serialized.get('id')
        elif hasattr(serialized, 'get'):
            model_name = serialized.get('name')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='llm_start',
            name=model_name,
            input_data={'prompts': prompts, 'serialized': serialized, 'tags': tags, 'metadata': metadata, 'kwargs': kwargs}
        )
    
    def on_chat_model_start(self, serialized, messages, run_id, parent_run_id=None, tags=None, metadata=None, **kwargs):
        # Extract model name from serialized
        model_name = None
        if isinstance(serialized, dict):
            model_name = serialized.get('name') or serialized.get('id')
        elif hasattr(serialized, 'get'):
            model_name = serialized.get('name')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='chat_model_start',
            name=model_name,
            input_data={'messages': messages, 'serialized': serialized, 'tags': tags, 'metadata': metadata, 'kwargs': kwargs}
        )
    
    def on_llm_new_token(self, token, run_id, parent_run_id=None, chunk=None, tags=None, **kwargs):
        # Note: We might want to skip logging every token to avoid too many DB writes
        # Uncomment if you want to log tokens
        # self._insert_log(
        #     run_id=run_id,
        #     parent_run_id=parent_run_id,
        #     event_type='llm_new_token',
        #     input_data={'token': token, 'chunk': str(chunk) if chunk else None, 'tags': tags, 'kwargs': kwargs}
        # )
    
    def on_llm_end(self, response, run_id, parent_run_id=None, tags=None, **kwargs):
        # Extract model name from response if available
        model_name = None
        if hasattr(response, 'llm_output') and response.llm_output:
            model_name = response.llm_output.get('model_name')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='llm_end',
            name=model_name,
            output_data={'response': str(response), 'tags': tags, 'kwargs': kwargs}
        )
    
    def on_llm_error(self, error, run_id, parent_run_id=None, tags=None, **kwargs):
        print(f"\033[91m🔴 LLM ERROR: {error}\033[0m")
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='llm_error',
            error=str(error),
            input_data={'tags': tags, 'kwargs': kwargs}
        )
    
    def on_chain_start(self, serialized, inputs, run_id, parent_run_id=None, tags=None, metadata=None, **kwargs):
        # Extract chain name from serialized
        chain_name = None
        if isinstance(serialized, dict):
            chain_name = serialized.get('name') or serialized.get('id')
        elif hasattr(serialized, 'get'):
            chain_name = serialized.get('name')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='chain_start',
            name=chain_name,
            input_data={'inputs': inputs, 'serialized': serialized, 'tags': tags, 'metadata': metadata, 'kwargs': kwargs}
        )
    
    def on_chain_end(self, outputs, run_id, parent_run_id=None, tags=None, **kwargs):
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='chain_end',
            output_data={'outputs': outputs, 'tags': tags, 'kwargs': kwargs}
        )
    
    def on_chain_error(self, error, run_id, parent_run_id=None, tags=None, **kwargs):
        print(f"\033[91m🔴 CHAIN ERROR: {error}\033[0m")
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='chain_error',
            error=str(error),
            input_data={'tags': tags, 'kwargs': kwargs}
        )
    
    def on_tool_start(self, serialized, input_str, run_id, parent_run_id=None, tags=None, metadata=None, inputs=None, **kwargs):
        # Extract tool name from serialized
        tool_name = None
        if isinstance(serialized, dict):
            tool_name = serialized.get('name') or serialized.get('id')
        elif hasattr(serialized, 'get'):
            tool_name = serialized.get('name')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='tool_start',
            name=tool_name,
            input_data={'input_str': input_str, 'inputs': inputs, 'serialized': serialized, 'tags': tags, 'metadata': metadata, 'kwargs': kwargs}
        )
    
    def on_tool_end(self, output, run_id, parent_run_id=None, **kwargs):
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='tool_end',
            output_data={'output': output, 'kwargs': kwargs}
        )
    
    def on_tool_error(self, error, run_id, parent_run_id=None, **kwargs):
        print(f"\033[91m🔴 TOOL ERROR: {error}\033[0m")
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='tool_error',
            error=str(error),
            input_data={'kwargs': kwargs}
        )
    
    def on_agent_action(self, action, run_id, parent_run_id=None, **kwargs):
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='agent_action',
            name=action.tool if hasattr(action, 'tool') else None,
            input_data={'tool': action.tool if hasattr(action, 'tool') else None, 
                       'tool_input': action.tool_input if hasattr(action, 'tool_input') else None,
                       'log': action.log if hasattr(action, 'log') else None,
                       'kwargs': kwargs}
        )
    
    def on_agent_finish(self, finish, run_id, parent_run_id=None, **kwargs):
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='agent_finish',
            output_data={'return_values': finish.return_values if hasattr(finish, 'return_values') else None,
                        'log': finish.log if hasattr(finish, 'log') else None,
                        'kwargs': kwargs}
        )
    
    def on_retriever_start(self, serialized, query, run_id, parent_run_id=None, tags=None, metadata=None, **kwargs):
        # Extract retriever name from serialized
        retriever_name = None
        if isinstance(serialized, dict):
            retriever_name = serialized.get('name') or serialized.get('id')
        elif hasattr(serialized, 'get'):
            retriever_name = serialized.get('name')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='retriever_start',
            name=retriever_name,
            input_data={'query': query, 'serialized': serialized, 'tags': tags, 'metadata': metadata, 'kwargs': kwargs}
        )
    
    def on_retriever_end(self, documents, run_id, parent_run_id=None, **kwargs):
        # Convert documents to serializable format
        docs_data = []
        if documents:
            for doc in documents:
                if hasattr(doc, 'page_content'):
                    docs_data.append({'page_content': doc.page_content, 'metadata': doc.metadata if hasattr(doc, 'metadata') else None})
                else:
                    docs_data.append(str(doc))
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='retriever_end',
            output_data={'documents': docs_data, 'document_count': len(documents) if documents else 0, 'kwargs': kwargs}
        )
    
    def on_retriever_error(self, error, run_id, parent_run_id=None, **kwargs):
        print(f"\033[91m🔴 RETRIEVER ERROR: {error}\033[0m")
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='retriever_error',
            error=str(error),
            input_data={'kwargs': kwargs}
        )
    
    def on_text(self, text, run_id, parent_run_id=None, **kwargs):
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='text',
            input_data={'text': text, 'kwargs': kwargs}
        )
    
    def on_retry(self, retry_state, run_id, parent_run_id=None, **kwargs):
        # Extract retry attempt info
        attempt = None
        if isinstance(retry_state, dict):
            attempt = retry_state.get('attempt')
        elif hasattr(retry_state, 'get'):
            attempt = retry_state.get('attempt')
        
        self._insert_log(
            run_id=run_id,
            parent_run_id=parent_run_id,
            event_type='retry',
            name=f"attempt_{attempt}" if attempt else None,
            input_data={'retry_state': str(retry_state), 'kwargs': kwargs}
        )
