"""
Structured logging for LLM requests and responses.

Logs prompts, tool calls, latency, and other metrics.
"""

import json
import logging
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, asdict

# Try to import jsonlogger, fall back to standard logging if not available
try:
    from python_json_logger import jsonlogger
    HAS_JSON_LOGGER = True
except ImportError:
    HAS_JSON_LOGGER = False

# Log directory
LOG_DIR = Path(__file__).parent.parent / "data" / "logs"
LOG_DIR.mkdir(parents=True, exist_ok=True)


class JSONFormatter(logging.Formatter):
    """Simple JSON formatter for logs."""
    
    def format(self, record):
        """Format log record as JSON."""
        log_data = {
            "timestamp": datetime.fromtimestamp(record.created).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage()
        }
        
        # Add extra fields
        if hasattr(record, '__dict__'):
            for key, value in record.__dict__.items():
                if key not in ['name', 'msg', 'args', 'created', 'filename', 
                              'funcName', 'levelname', 'levelno', 'lineno', 
                              'module', 'msecs', 'message', 'pathname', 'process',
                              'processName', 'relativeCreated', 'thread', 'threadName',
                              'exc_info', 'exc_text', 'stack_info']:
                    log_data[key] = value
        
        return json.dumps(log_data)


@dataclass
class LLMRequestLog:
    """Structured log entry for an LLM request."""
    timestamp: str
    session_id: Optional[str]
    agent_type: str  # "work" or "family"
    user_id: Optional[str]
    request_id: str
    prompt: str
    messages_count: int
    tools_available: int
    tools_called: Optional[List[str]]
    latency_ms: float
    tokens_in: Optional[int]
    tokens_out: Optional[int]
    response_length: int
    error: Optional[str]
    model: str


class LLMLogger:
    """Structured logger for LLM operations."""
    
    def __init__(self, log_file: Optional[Path] = None):
        """Initialize logger."""
        if log_file is None:
            log_file = LOG_DIR / f"llm_{datetime.now().strftime('%Y%m%d')}.log"
        
        self.log_file = log_file
        
        # Set up JSON logger
        self.logger = logging.getLogger("llm")
        self.logger.setLevel(logging.INFO)
        
        # File handler with JSON formatter
        file_handler = logging.FileHandler(str(log_file))
        if HAS_JSON_LOGGER:
            formatter = jsonlogger.JsonFormatter()
        else:
            # Fallback: Use custom JSON formatter
            formatter = JSONFormatter()
        file_handler.setFormatter(formatter)
        self.logger.addHandler(file_handler)
        
        # Also log to console (structured)
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        self.logger.addHandler(console_handler)
    
    def log_request(self, 
                   session_id: Optional[str],
                   agent_type: str,
                   user_id: Optional[str],
                   request_id: str,
                   prompt: str,
                   messages: List[Dict[str, Any]],
                   tools_available: int,
                   start_time: float,
                   end_time: float,
                   response: Dict[str, Any],
                   tools_called: Optional[List[str]] = None,
                   error: Optional[str] = None,
                   model: str = "unknown"):
        """Log an LLM request."""
        latency_ms = (end_time - start_time) * 1000
        
        # Extract token counts if available
        tokens_in = response.get("prompt_eval_count")
        tokens_out = response.get("eval_count")
        response_text = response.get("message", {}).get("content", "")
        
        log_entry = LLMRequestLog(
            timestamp=datetime.now().isoformat(),
            session_id=session_id,
            agent_type=agent_type,
            user_id=user_id,
            request_id=request_id,
            prompt=prompt[:500],  # Truncate long prompts
            messages_count=len(messages),
            tools_available=tools_available,
            tools_called=tools_called or [],
            latency_ms=round(latency_ms, 2),
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            response_length=len(response_text),
            error=error,
            model=model
        )
        
        # Log as JSON
        self.logger.info("LLM Request", extra=asdict(log_entry))
    
    def log_error(self,
                 session_id: Optional[str],
                 agent_type: str,
                 request_id: str,
                 error: str,
                 context: Optional[Dict[str, Any]] = None):
        """Log an error."""
        log_data = {
            "timestamp": datetime.now().isoformat(),
            "type": "error",
            "session_id": session_id,
            "agent_type": agent_type,
            "request_id": request_id,
            "error": error
        }
        
        if context:
            log_data.update(context)
        
        self.logger.error("LLM Error", extra=log_data)


# Global logger instance
_logger = LLMLogger()


def get_llm_logger() -> LLMLogger:
    """Get the global LLM logger instance."""
    return _logger