"""
Core logger implementation for LLM training data collection.
"""

import atexit
import json
import queue
import threading
import time
from typing import Any, Dict, Optional
import logging

try:
    import requests
except ImportError:
    raise ImportError(
        "The 'requests' library is required. Install it with: pip install requests"
    )


class LLMLogger:
    """
    A thread-safe, asynchronous logger for LLM request/response data.
    
    Logs are batched and sent to a backend API in the background without
    blocking the main application. Includes automatic retry with exponential
    backoff and graceful shutdown handling.
    
    Args:
        endpoint (str): The backend API endpoint to send logs to
        api_key (str): API key for authentication
        batch_size (int): Number of logs to batch before sending (default: 10)
        flush_interval (float): Seconds to wait before flushing incomplete batch (default: 5.0)
        max_queue_size (int): Maximum number of logs to queue (default: 10000)
        max_retries (int): Maximum number of retry attempts (default: 3)
        initial_retry_delay (float): Initial delay for exponential backoff in seconds (default: 1.0)
        max_retry_delay (float): Maximum retry delay in seconds (default: 60.0)
        timeout (float): HTTP request timeout in seconds (default: 10.0)
        debug (bool): Enable debug logging (default: False)
    """
    
    def __init__(
        self,
        endpoint: str,
        api_key: str,
        batch_size: int = 10,
        flush_interval: float = 5.0,
        max_queue_size: int = 10000,
        max_retries: int = 3,
        initial_retry_delay: float = 1.0,
        max_retry_delay: float = 60.0,
        timeout: float = 10.0,
        debug: bool = False,
    ):
        self.endpoint = endpoint
        self.api_key = api_key
        self.batch_size = batch_size
        self.flush_interval = flush_interval
        self.max_retries = max_retries
        self.initial_retry_delay = initial_retry_delay
        self.max_retry_delay = max_retry_delay
        self.timeout = timeout
        
        # Setup logging
        self.logger = logging.getLogger(__name__)
        if debug:
            self.logger.setLevel(logging.DEBUG)
            if not self.logger.handlers:
                handler = logging.StreamHandler()
                handler.setFormatter(
                    logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
                )
                self.logger.addHandler(handler)
        
        # Thread-safe queue for log entries
        self.log_queue: queue.Queue = queue.Queue(maxsize=max_queue_size)
        
        # Control flags
        self._shutdown = threading.Event()
        self._worker_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()
        
        # Statistics
        self.stats = {
            "logs_queued": 0,
            "logs_sent": 0,
            "logs_failed": 0,
            "batches_sent": 0,
            "batches_failed": 0,
        }
        
        # Start background worker thread
        self._start_worker()
        
        # Register cleanup on program exit
        atexit.register(self.shutdown)
    
    def _start_worker(self):
        """Start the background worker thread."""
        with self._lock:
            if self._worker_thread is None or not self._worker_thread.is_alive():
                self._shutdown.clear()
                self._worker_thread = threading.Thread(
                    target=self._worker_loop,
                    daemon=True,
                    name="LLMLogger-Worker"
                )
                self._worker_thread.start()
                self.logger.debug("Worker thread started")
    
    def _worker_loop(self):
        """
        Main worker loop that batches and sends logs.
        Runs in background thread until shutdown signal received.
        """
        batch = []
        last_flush_time = time.time()
        
        while not self._shutdown.is_set():
            try:
                # Try to get a log entry with timeout
                try:
                    log_entry = self.log_queue.get(timeout=0.5)
                    batch.append(log_entry)
                    self.log_queue.task_done()
                except queue.Empty:
                    pass
                
                # Check if we should flush the batch
                should_flush = (
                    len(batch) >= self.batch_size or
                    (batch and time.time() - last_flush_time >= self.flush_interval)
                )
                
                if should_flush and batch:
                    self._send_batch(batch)
                    batch = []
                    last_flush_time = time.time()
                    
            except Exception as e:
                self.logger.error(f"Error in worker loop: {e}", exc_info=True)
        
        # Flush remaining logs on shutdown
        if batch:
            self.logger.debug(f"Flushing {len(batch)} remaining logs on shutdown")
            self._send_batch(batch)
    
    def _send_batch(self, batch: list):
        """
        Send a batch of logs to the backend API with retry logic.
        
        Args:
            batch (list): List of log entries to send
        """
        if not batch:
            return
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        
        payload = {
            "logs": batch,
        }
        
        retry_delay = self.initial_retry_delay
        
        for attempt in range(self.max_retries + 1):
            try:
                self.logger.debug(
                    f"Sending batch of {len(batch)} logs (attempt {attempt + 1}/{self.max_retries + 1})"
                )
                
                response = requests.post(
                    self.endpoint,
                    json=payload,
                    headers=headers,
                    timeout=self.timeout,
                )
                
                if response.status_code in (200, 201, 202):
                    # Success
                    self.stats["logs_sent"] += len(batch)
                    self.stats["batches_sent"] += 1
                    self.logger.debug(
                        f"Successfully sent batch of {len(batch)} logs. "
                        f"Status: {response.status_code}"
                    )
                    return
                elif response.status_code >= 500:
                    # Server error - retry
                    self.logger.warning(
                        f"Server error {response.status_code}. "
                        f"Attempt {attempt + 1}/{self.max_retries + 1}"
                    )
                elif response.status_code == 429:
                    # Rate limit - retry
                    self.logger.warning(
                        f"Rate limited. Attempt {attempt + 1}/{self.max_retries + 1}"
                    )
                else:
                    # Client error (4xx) - don't retry
                    self.logger.error(
                        f"Client error {response.status_code}: {response.text}. "
                        "Not retrying."
                    )
                    self.stats["logs_failed"] += len(batch)
                    self.stats["batches_failed"] += 1
                    return
                    
            except requests.exceptions.Timeout:
                self.logger.warning(
                    f"Request timeout. Attempt {attempt + 1}/{self.max_retries + 1}"
                )
            except requests.exceptions.ConnectionError:
                self.logger.warning(
                    f"Connection error. Attempt {attempt + 1}/{self.max_retries + 1}"
                )
            except Exception as e:
                self.logger.error(
                    f"Unexpected error sending batch: {e}. "
                    f"Attempt {attempt + 1}/{self.max_retries + 1}",
                    exc_info=True
                )
            
            # If we haven't exhausted retries, wait before retrying
            if attempt < self.max_retries:
                time.sleep(retry_delay)
                # Exponential backoff with max cap
                retry_delay = min(retry_delay * 2, self.max_retry_delay)
        
        # All retries exhausted
        self.logger.error(
            f"Failed to send batch of {len(batch)} logs after {self.max_retries + 1} attempts. "
            "Dropping logs."
        )
        self.stats["logs_failed"] += len(batch)
        self.stats["batches_failed"] += 1
    
    def log(
        self,
        target_model: str,
        request: Dict[str, Any],
        response: Dict[str, Any],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> bool:
        """
        Log an LLM request/response pair.
        
        This method returns immediately and queues the log for asynchronous processing.
        
        Args:
            request (dict): The LLM request data (prompt, model, parameters, etc.)
            response (dict): The LLM response data (completion, tokens, etc.)
            metadata (dict, optional): Additional metadata (user_id, session_id, etc.)
        
        Returns:
            bool: True if log was queued successfully, False if queue is full
        """
        log_entry = {
            "target_model": target_model,
            "request": request,
            "response": response,
            "metadata": metadata or {},
            "timestamp": time.time(),
        }
        
        try:
            # Non-blocking put with immediate timeout
            self.log_queue.put_nowait(log_entry)
            self.stats["logs_queued"] += 1
            self.logger.debug("Log entry queued successfully")
            return True
        except queue.Full:
            self.logger.warning(
                "Log queue is full. Dropping log entry. "
                f"Current queue size: {self.log_queue.qsize()}"
            )
            self.stats["logs_failed"] += 1
            return False
    
    def flush(self, timeout: Optional[float] = None):
        """
        Flush all queued logs and wait for them to be sent.
        
        Args:
            timeout (float, optional): Maximum time to wait in seconds. 
                                      None means wait indefinitely.
        """
        self.logger.debug(f"Flushing queue. Current size: {self.log_queue.qsize()}")
        
        # Wait for queue to be empty
        if timeout is None:
            self.log_queue.join()
        else:
            start_time = time.time()
            while not self.log_queue.empty() and time.time() - start_time < timeout:
                time.sleep(0.1)
        
        self.logger.debug("Flush complete")
    
    def shutdown(self, timeout: float = 10.0):
        """
        Gracefully shutdown the logger, flushing all remaining logs.
        
        Args:
            timeout (float): Maximum time to wait for shutdown in seconds
        """
        if self._shutdown.is_set():
            return  # Already shut down
        
        self.logger.debug("Shutting down logger...")
        
        # Signal shutdown and wait for worker to finish
        self._shutdown.set()
        
        if self._worker_thread and self._worker_thread.is_alive():
            self._worker_thread.join(timeout=timeout)
            
            if self._worker_thread.is_alive():
                self.logger.warning(
                    f"Worker thread did not terminate within {timeout}s timeout"
                )
        
        self.logger.debug("Logger shutdown complete")
    
    def get_stats(self) -> Dict[str, int]:
        """
        Get statistics about logger performance.
        
        Returns:
            dict: Dictionary containing statistics
        """
        return self.stats.copy()
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - ensure graceful shutdown."""
        self.shutdown()
        return False
    
    def __del__(self):
        """Destructor - ensure cleanup."""
        try:
            self.shutdown(timeout=2.0)
        except Exception:
            pass  # Ignore errors during cleanup

