"""
Rollout Logging Module

This module provides comprehensive logging functionality for agent rollout system traces,
including timing, performance metrics, and detailed analysis capabilities.
"""

import time
import logging
import csv
import os
import json
import numpy as np
from datetime import datetime
from dataclasses import dataclass
from typing import Dict, Any, Optional, List

# Configure logging
logger = logging.getLogger(__name__)

@dataclass
class RolloutTrace:
    """Data structure for agent rollout system traces."""
    task_id: str
    agent_type: str
    rollout_id: str
    start_time: float
    end_time: Optional[float] = None
    duration: Optional[float] = None
    num_turns: int = 0
    input_length: int = 0
    output_length: int = 0
    input_token_count: int = 0
    output_token_count: int = 0
    final_reward: Optional[float] = None
    status: str = "running"
    error_message: Optional[str] = None
    model: Optional[str] = None
    endpoint: Optional[str] = None
    retry_count: int = 0
    max_retries: int = 0
    timeout_seconds: float = 60.0
    is_retry: bool = False
    original_task_id: Optional[str] = None
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}

class RolloutLogger:
    """Handles logging of agent rollout system traces."""
    
    def __init__(self, log_dir: str = "logs", experiment_name: str = "agentarium"):
        """
        Initialize rollout logger.
        
        Args:
            log_dir: Directory to store log files
            experiment_name: Name of the experiment for the log filename
        """
        self.log_dir = log_dir
        self.experiment_name = experiment_name
        self.traces: Dict[str, RolloutTrace] = {}
        
        # Create log directory
        os.makedirs(log_dir, exist_ok=True)
        
        # Create CSV file path with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_filename = f"{experiment_name}_rollout_traces_{timestamp}.csv"
        self.csv_filepath = os.path.join(log_dir, csv_filename)
        self.csv_headers_written = False
        
        logger.info(f"RolloutLogger initialized. Traces will be saved to: {self.csv_filepath}")
    
    def start_rollout_trace(self, task_id: str, agent_type: str, rollout_id: str, 
                           model: str = None, endpoint: str = None, metadata: Dict[str, Any] = None,
                           retry_count: int = 0, max_retries: int = 0, timeout_seconds: float = 60.0,
                           is_retry: bool = False, original_task_id: str = None) -> RolloutTrace:
        """
        Start a new rollout trace.
        
        Args:
            task_id: Unique task identifier
            agent_type: Type of agent being used
            rollout_id: Unique rollout identifier
            model: Model name being used
            endpoint: Endpoint being used
            metadata: Additional metadata
            retry_count: Current retry attempt number
            max_retries: Maximum number of retries allowed
            timeout_seconds: Timeout duration for this task
            is_retry: Whether this is a retry attempt
            original_task_id: Original task ID if this is a retry
            
        Returns:
            RolloutTrace object
        """
        trace = RolloutTrace(
            task_id=task_id,
            agent_type=agent_type,
            rollout_id=rollout_id,
            start_time=time.time(),
            model=model,
            endpoint=endpoint,
            retry_count=retry_count,
            max_retries=max_retries,
            timeout_seconds=timeout_seconds,
            is_retry=is_retry,
            original_task_id=original_task_id,
            metadata=metadata or {}
        )
        
        self.traces[task_id] = trace
        retry_info = f" (retry {retry_count}/{max_retries})" if is_retry else ""
        logger.debug(f"Started rollout trace for task {task_id}{retry_info}")
        return trace
    
    def complete_rollout_trace(self, task_id: str, result: Any = None, error: str = None):
        """
        Complete a rollout trace with results.
        
        Args:
            task_id: Task identifier
            result: Rollout result object
            error: Error message if the rollout failed
        """
        if task_id not in self.traces:
            logger.warning(f"No trace found for task {task_id}")
            return
            
        trace = self.traces[task_id]
        trace.end_time = time.time()
        trace.duration = trace.end_time - trace.start_time
        
        if error:
            trace.status = "failed"
            trace.error_message = error
        else:
            trace.status = "completed"
            
            # Extract metrics from result if available
            if result and hasattr(result, 'final_reward'):
                trace.final_reward = result.final_reward
                
            if result and hasattr(result, 'triplets') and result.triplets:
                # FIXME
                trace.num_turns = len(result.triplets)
                
                # Calculate input/output lengths
                for triplet in result.triplets:
                    if hasattr(triplet, 'prompt'):
                        if isinstance(triplet.prompt, dict):
                            prompt_text = triplet.prompt.get('text', '')
                            trace.input_length += len(prompt_text)
                            prompt_tokens = triplet.prompt.get('token_ids', [])
                            trace.input_token_count += len(prompt_tokens)
                        elif isinstance(triplet.prompt, str):
                            trace.input_length += len(triplet.prompt)
                            
                    if hasattr(triplet, 'response'):
                        if isinstance(triplet.response, dict):
                            response_text = triplet.response.get('text', '')
                            trace.output_length += len(response_text)
                            response_tokens = triplet.response.get('token_ids', [])
                            trace.output_token_count += len(response_tokens)
                        elif isinstance(triplet.response, str):
                            trace.output_length += len(triplet.response)
        
        # Log the completed trace to CSV
        self._log_trace_to_csv(trace)
        
        logger.debug(f"Completed rollout trace for task {task_id}: "
                    f"duration={trace.duration:.2f}s, turns={trace.num_turns}, "
                    f"status={trace.status}")
    
    def get_trace(self, task_id: str) -> Optional[RolloutTrace]:
        """Get a trace by task ID."""
        return self.traces.get(task_id)
    
    def get_all_traces(self) -> Dict[str, RolloutTrace]:
        """Get all traces."""
        return self.traces.copy()
    
    def get_completed_traces(self) -> List[RolloutTrace]:
        """Get all completed traces."""
        return [trace for trace in self.traces.values() if trace.status == "completed"]
    
    def get_summary_stats(self) -> Dict[str, Any]:
        """Get summary statistics of all traces."""
        completed_traces = self.get_completed_traces()
        
        if not completed_traces:
            return {
                "total_traces": len(self.traces),
                "completed_traces": 0,
                "failed_traces": len([t for t in self.traces.values() if t.status == "failed"]),
                "running_traces": len([t for t in self.traces.values() if t.status == "running"])
            }
        
        durations = [t.duration for t in completed_traces if t.duration is not None]
        rewards = [t.final_reward for t in completed_traces if t.final_reward is not None]
        turns = [t.num_turns for t in completed_traces]
        input_lengths = [t.input_length for t in completed_traces]
        output_lengths = [t.output_length for t in completed_traces]
        
        return {
            "total_traces": len(self.traces),
            "completed_traces": len(completed_traces),
            "failed_traces": len([t for t in self.traces.values() if t.status == "failed"]),
            "running_traces": len([t for t in self.traces.values() if t.status == "running"]),
            "avg_duration": np.mean(durations) if durations else 0,
            "min_duration": np.min(durations) if durations else 0,
            "max_duration": np.max(durations) if durations else 0,
            "avg_reward": np.mean(rewards) if rewards else 0,
            "avg_turns": np.mean(turns) if turns else 0,
            "avg_input_length": np.mean(input_lengths) if input_lengths else 0,
            "avg_output_length": np.mean(output_lengths) if output_lengths else 0,
            "total_input_length": sum(input_lengths),
            "total_output_length": sum(output_lengths)
        }
    
    def _log_trace_to_csv(self, trace: RolloutTrace):
        """Log a trace to the CSV file."""
        try:
            # Convert trace to dictionary
            trace_dict = {
                'task_id': trace.task_id,
                'agent_type': trace.agent_type,
                'rollout_id': trace.rollout_id,
                'start_time': trace.start_time,
                'end_time': trace.end_time,
                'duration': trace.duration,
                'num_turns': trace.num_turns,
                'input_length': trace.input_length,
                'output_length': trace.output_length,
                'input_token_count': trace.input_token_count,
                'output_token_count': trace.output_token_count,
                'final_reward': trace.final_reward,
                'status': trace.status,
                'error_message': trace.error_message,
                'model': trace.model,
                'endpoint': trace.endpoint,
                'retry_count': trace.retry_count,
                'max_retries': trace.max_retries,
                'timeout_seconds': trace.timeout_seconds,
                'is_retry': trace.is_retry,
                'original_task_id': trace.original_task_id,
                'metadata': json.dumps(trace.metadata) if trace.metadata else ""
            }
            
            # Write to CSV
            file_exists = os.path.exists(self.csv_filepath)
            with open(self.csv_filepath, 'a', newline='') as csvfile:
                fieldnames = list(trace_dict.keys())
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                
                if not file_exists or not self.csv_headers_written:
                    writer.writeheader()
                    self.csv_headers_written = True
                
                writer.writerow(trace_dict)
                
        except Exception as e:
            logger.warning(f"Failed to log trace to CSV: {e}")
    
    @property
    def filepath(self) -> str:
        """Get the CSV file path."""
        return self.csv_filepath
