"""Enhanced JSON logging system for agent evaluations."""

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


@dataclass
class LogEntry:
    """Represents a single log entry."""
    timestamp: str
    level: str
    message: str
    data: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return asdict(self)


class JSONLogger:
    """Enhanced JSON logger for evaluation runs."""
    
    def __init__(
        self,
        output_dir: str,
        experiment_name: Optional[str] = None,
        logger: Optional[logging.Logger] = None
    ):
        """Initialize JSON logger.
        
        Args:
            output_dir: Directory for log files
            experiment_name: Name of the experiment
            logger: Python logger instance
        """
        self.output_dir = Path(output_dir).expanduser()
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        self.experiment_name = experiment_name or datetime.now().strftime("%Y%m%d_%H%M%S")
        self.logger = logger or logging.getLogger(__name__)
        
        # Create log file
        self.log_file = self.output_dir / f"{self.experiment_name}_log.json"
        self.summary_file = self.output_dir / f"{self.experiment_name}_summary.json"
        
        # In-memory log storage
        self.entries: List[LogEntry] = []
        self.metadata: Dict[str, Any] = {
            'experiment_name': self.experiment_name,
            'start_time': datetime.now().isoformat(),
            'output_dir': str(self.output_dir)
        }
    
    def log_evaluation(
        self,
        task_id: str,
        model_name: str,
        problem_data: Dict[str, Any],
        response: Dict[str, Any],
        metrics: Dict[str, Any],
        config: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Log a complete evaluation run.
        
        Args:
            task_id: Unique task identifier
            model_name: Name of the model
            problem_data: Problem information
            response: Agent response data
            metrics: Evaluation metrics
            config: Configuration used
            
        Returns:
            Complete log entry dictionary
        """
        # Extract key information
        log_data = {
            'task_id': task_id,
            'timestamp': datetime.now().isoformat(),
            'model': {
                'name': model_name,
                'temperature': config.get('temperature', 'N/A'),
                'provider': config.get('provider', 'N/A')
            },
            'problem': {
                'id': problem_data.get('question_id', 'unknown'),
                'question': problem_data.get('question', ''),
                'ground_truth': problem_data.get('final_answer', '')
            },
            'response': self._process_response(response),
            'metrics': metrics,
            'config_hash': self._hash_config(config)
        }
        
        # Create log entry
        entry = LogEntry(
            timestamp=log_data['timestamp'],
            level='INFO',
            message=f"Evaluation completed for {task_id}",
            data=log_data
        )
        
        self.entries.append(entry)
        
        # Save to file
        self._save_entry(entry)
        
        return log_data
    
    def _process_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
        """Process agent response for logging.
        
        Args:
            response: Raw agent response
            
        Returns:
            Processed response data
        """
        processed = {
            'final_answer': None,
            'conversation_length': 0,
            'tool_calls': [],
            'total_tokens': 0,
            'full_conversation': []  # Add full conversation history
        }
        
        # Extract messages
        messages = response.get('messages', [])
        processed['conversation_length'] = len(messages)
        
        # Save full conversation history
        for msg in messages:
            msg_data = {
                'role': getattr(msg, 'role', 'unknown') if hasattr(msg, 'role') else str(type(msg).__name__),
                'content': getattr(msg, 'content', '') if hasattr(msg, 'content') else str(msg),
            }
            
            # Add tool calls if present
            if hasattr(msg, 'tool_calls') and msg.tool_calls:
                msg_data['tool_calls'] = [
                    {
                        'name': tc.get('name', 'unknown') if isinstance(tc, dict) else getattr(tc, 'name', 'unknown'),
                        'args': tc.get('args', {}) if isinstance(tc, dict) else getattr(tc, 'args', {})
                    }
                    for tc in msg.tool_calls
                ]
            
            # Add additional metadata if present
            if hasattr(msg, 'response_metadata'):
                msg_data['metadata'] = msg.response_metadata
                
            processed['full_conversation'].append(msg_data)
        
        # Extract final answer
        if messages:
            last_message = messages[-1]
            if hasattr(last_message, 'content'):
                content = last_message.content
                # Look for answer pattern
                if isinstance(content, str) and '####' in content:
                    parts = content.split('####')
                    if len(parts) > 1:
                        processed['final_answer'] = parts[1].strip().split()[0]
        
        # Extract tool calls
        for msg in messages:
            if hasattr(msg, 'tool_calls') and msg.tool_calls:
                for tool_call in msg.tool_calls:
                    processed['tool_calls'].append({
                        'name': tool_call.get('name', 'unknown'),
                        'args': tool_call.get('args', {})
                    })
        
        # Extract token usage
        if messages and hasattr(messages[-1], 'response_metadata'):
            metadata = messages[-1].response_metadata
            if metadata and 'usage' in metadata:
                usage = metadata['usage']
                processed['total_tokens'] = usage.get('input_tokens', 0) + usage.get('output_tokens', 0)
        
        return processed
    
    def _hash_config(self, config: Dict[str, Any]) -> str:
        """Generate hash of configuration for tracking.
        
        Args:
            config: Configuration dictionary
            
        Returns:
            SHA256 hash of configuration
        """
        config_str = json.dumps(config, sort_keys=True)
        return hashlib.sha256(config_str.encode()).hexdigest()[:16]
    
    def _save_entry(self, entry: LogEntry):
        """Save a log entry to file.
        
        Args:
            entry: Log entry to save
        """
        try:
            # Append to log file
            with open(self.log_file, 'a') as f:
                json.dump(entry.to_dict(), f)
                f.write('\n')
        except Exception as e:
            self.logger.error(f"Failed to save log entry: {e}")
    
    def save_summary(self, summary_data: Dict[str, Any]):
        """Save evaluation summary.
        
        Args:
            summary_data: Summary statistics and results
        """
        self.metadata['end_time'] = datetime.now().isoformat()
        self.metadata['total_entries'] = len(self.entries)
        
        summary = {
            'metadata': self.metadata,
            'summary': summary_data,
            'entries_count': len(self.entries)
        }
        
        try:
            with open(self.summary_file, 'w') as f:
                json.dump(summary, f, indent=2)
            self.logger.info(f"Saved summary to {self.summary_file}")
        except Exception as e:
            self.logger.error(f"Failed to save summary: {e}")
    
    def get_entries_by_model(self, model_name: str) -> List[LogEntry]:
        """Get all log entries for a specific model.
        
        Args:
            model_name: Name of the model
            
        Returns:
            List of log entries for the model
        """
        return [
            entry for entry in self.entries
            if entry.data.get('model', {}).get('name') == model_name
        ]
    
    def get_entries_by_problem(self, problem_id: str) -> List[LogEntry]:
        """Get all log entries for a specific problem.
        
        Args:
            problem_id: ID of the problem
            
        Returns:
            List of log entries for the problem
        """
        return [
            entry for entry in self.entries
            if entry.data.get('problem', {}).get('id') == problem_id
        ]
    
    def export_to_csv(self, output_file: Optional[str] = None):
        """Export log entries to CSV format.
        
        Args:
            output_file: Path to output CSV file
        """
        import pandas as pd
        
        if not output_file:
            output_file = self.output_dir / f"{self.experiment_name}_log.csv"
        
        # Flatten log entries for CSV
        rows = []
        for entry in self.entries:
            row = {
                'timestamp': entry.timestamp,
                'task_id': entry.data.get('task_id'),
                'model_name': entry.data.get('model', {}).get('name'),
                'problem_id': entry.data.get('problem', {}).get('id'),
                'ground_truth': entry.data.get('problem', {}).get('ground_truth'),
                'final_answer': entry.data.get('response', {}).get('final_answer'),
                'is_correct': entry.data.get('metrics', {}).get('is_correct'),
                'total_tokens': entry.data.get('response', {}).get('total_tokens'),
                'conversation_length': entry.data.get('response', {}).get('conversation_length')
            }
            rows.append(row)
        
        df = pd.DataFrame(rows)
        df.to_csv(output_file, index=False)
        self.logger.info(f"Exported {len(rows)} entries to {output_file}")
