import os
from datetime import datetime
from typing import TextIO, Optional

class Colors:
    GREEN = "\033[92m"
    RED = "\033[91m"
    YELLOW = "\033[93m"
    BLUE = "\033[94m"
    MAGENTA = "\033[95m"
    CYAN = "\033[96m"
    RESET = "\033[0m"

class NesyLogger:
    """Logger class specifically for NeSy Agent with enhanced functionality"""
    
    def __init__(self, log_dir: str, debug: bool = False) -> None:
        self.log_dir: str = log_dir
        self.debug_mode: bool = debug
        self.timestamp: str = self._get_timestamp()
        
        # Create logging directory
        os.makedirs(log_dir, exist_ok=True)
        
        # Initialize file streams
        log_filename: str = f"nesy_agent_{self.timestamp}.log"
        error_filename: str = f"nesy_agent_{self.timestamp}.error"
        debug_filename: str = f"nesy_agent_{self.timestamp}.debug"
        
        log_filepath: str = os.path.join(self.log_dir, log_filename)
        error_filepath: str = os.path.join(self.log_dir, error_filename)
        debug_filepath: str = os.path.join(self.log_dir, debug_filename)
        
        self.log_file: TextIO = open(log_filepath, "w", encoding="utf-8")
        self.error_file: TextIO = open(error_filepath, "w", encoding="utf-8")
        self.debug_file: TextIO = open(debug_filepath, "w", encoding="utf-8")

    def __del__(self) -> None:
        """Close all file streams when logger is destroyed"""
        if hasattr(self, 'log_file'):
            self.log_file.close()
        if hasattr(self, 'error_file'):
            self.error_file.close()
        if hasattr(self, 'debug_file'):
            self.debug_file.close()
    
    def _get_timestamp(self) -> str:
        """Get current timestamp in YYYYMMDD_HHMMSS format"""
        return datetime.now().strftime("%Y%m%d_%H%M%S")
    
    def _write_to_stream(self, file_obj: TextIO, message: str, level: str = "INFO") -> None:
        """Write message to file stream with timestamp"""
        timestamp: str = self._get_timestamp()
        file_obj.write(f"[{timestamp}] [{level}] {message}\n")
        file_obj.flush()
    
    def info(self, message: str) -> None:
        """Log info message"""
        self._write_to_stream(self.log_file, message, "INFO")
        if self.debug_mode:
            print(f"{Colors.GREEN}[INFO]{Colors.RESET} {message}")
    
    def error(self, message: str) -> None:
        """Log error message"""
        self._write_to_stream(self.error_file, message, "ERROR")
        print(f"{Colors.RED}[ERROR]{Colors.RESET} {message}")
    
    def warning(self, message: str) -> None:
        """Log warning message"""
        self._write_to_stream(self.log_file, message, "WARNING")
        if self.debug_mode:
            print(f"{Colors.YELLOW}[WARNING]{Colors.RESET} {message}")
    
    def debug(self, message: str) -> None:
        """Log debug message"""
        self._write_to_stream(self.debug_file, message, "DEBUG")
        if self.debug_mode:
            print(f"{Colors.CYAN}[DEBUG]{Colors.RESET} {message}")
    
    def step(self, message: str) -> None:
        """Log step message for tracking execution flow"""
        self._write_to_stream(self.log_file, message, "STEP")
        if self.debug_mode:
            print(f"{Colors.BLUE}[STEP]{Colors.RESET} {message}")
    
    def constraint(self, message: str) -> None:
        """Log constraint evaluation message"""
        self._write_to_stream(self.log_file, message, "CONSTRAINT")
        if self.debug_mode:
            print(f"{Colors.MAGENTA}[CONSTRAINT]{Colors.RESET} {message}")
    
    def search(self, message: str) -> None:
        """Log search-related message"""
        self._write_to_stream(self.log_file, message, "SEARCH")
        if self.debug_mode:
            print(f"{Colors.CYAN}[SEARCH]{Colors.RESET} {message}")
    
    def performance(self, message: str) -> None:
        """Log performance-related message"""
        self._write_to_stream(self.log_file, message, "PERFORMANCE")
        if self.debug_mode:
            print(f"{Colors.YELLOW}[PERFORMANCE]{Colors.RESET} {message}")
