# agent.py
from typing import Any, Dict, Optional, List
from .config import Config
from agentforge.apis.base_api import BaseModel
from agentforge.utils.logger import Logger
from agentforge.utils.prompt_processor import PromptProcessor
from agentforge.utils.parsing_processor import ParsingProcessor
from agentforge.utils.audio_manager import AudioManager


class Agent:
    def __init__(self, agent_name: Optional[str] = None, log_file: Optional[str] = 'agentforge'):
        """
        Initialize an Agent instance with necessary configurations and services.
        
        Args:
            agent_name: Name of the agent. Defaults to class name if not provided.
            log_file: Name of the log file. Defaults to 'agentforge'.
        """
        self.agent_name: str = agent_name if agent_name is not None else self.__class__.__name__
        self.logger: Logger = Logger(self.agent_name, log_file)

        self.prompt = None
        self.result = None
        self.output = None
        self.parsed_result = None
        self.persona = None
        self.images = []

        # Initialize services
        self.config = Config()
        self.prompt_processor = PromptProcessor()
        self.parsing_processor = ParsingProcessor()
        
        # Initialize attributes and configurations
        self._initialize_data_attributes()
        self._initialize_agent_config()

    def _initialize_data_attributes(self) -> None:
        """Initialize all agent data attributes to their default values."""
        # Configuration
        self.agent_config: Optional[Any] = None
        self.model: Optional[BaseModel] = None
        self.persona: Optional[Dict[str, Any]] = None
        
        # Prompt related
        self.prompt: Optional[Dict[str]] = None
        self.prompt_template: Optional[Dict[str, Any]] = None
        self.template_data: Dict[str, Any] = {}
        
        # Results and output
        self.result: Optional[str] = None
        self.parsed_result: Optional[Any] = None
        self.output: Optional[str] = None

        # Media
        self.images: List[str] = []
        
        # Audio
        self.audio_manager: AudioManager | None = None

    # ---------------------------------
    # Execution
    # ---------------------------------

    def run(self, **kwargs: Any) -> Optional[str]:
        """
        Execute the agent's workflow from data loading through output building.
        
        Args:
            **kwargs: Keyword arguments to incorporate into the agent's data.
            
        Returns:
            The output generated by the agent or None if execution failed.
        """
        try:
            self.logger.info(f"{self.agent_name} - Running...")
            self._execute_workflow(**kwargs)
            self.logger.info(f"{self.agent_name} - Done!")
            return self.output
        except Exception as e:
            self.logger.error(f"Agent execution failed: {e}")
            return None

    def _execute_workflow(self, **kwargs: Any) -> None:
        """Execute the complete agent workflow steps."""
        self.load_data(**kwargs)
        self.process_data()
        self.render_prompt()
        self.run_model()
        self.parse_result()
        self.post_process_result()
        self.build_output()

    # ---------------------------------
    # Configuration Loading
    # ---------------------------------

    def _initialize_agent_config(self) -> None:
        """Load all agent configurations."""
        self.agent_config = self.config.load_agent_data(self.agent_name)
        self.prompt_template = self.agent_config.prompts
        self.model = self.agent_config.model
        self.load_persona_data()
        # Instantiate AudioManager after config is ready
        self.audio_manager = AudioManager(self.agent_config, self.logger)
        
    def load_persona_data(self) -> None:
        """Load persona data if personas are enabled in system settings."""
        if not self.agent_config.settings.system.persona.enabled or not self.agent_config.persona:
            return
        
        self.persona = self.agent_config.persona.copy()
        self.template_data['persona'] = self.persona
        self.logger.debug(f"Persona Data Loaded for '{self.agent_name}'.")

    # ---------------------------------
    # Data Loading
    # ---------------------------------

    def load_data(self, **kwargs: Any) -> None:
        """
        Load all data needed for prompt generation, including dynamic updates.
        
        Args:
            **kwargs: Additional data to incorporate into template variables.
        """
        if self.agent_config.settings.system.misc.on_the_fly:
            self._initialize_agent_config()

        self.load_additional_data()
        self.template_data.update(kwargs)

    def load_additional_data(self) -> None:
        """
        Load custom additional data for the agent.
        Override this method in subclasses to load custom data.
        """
        pass

    # ---------------------------------
    # Processing
    # ---------------------------------

    def process_data(self) -> None:
        """
        Process loaded data before generating prompts.
        Override this method in subclasses to implement custom data processing.
        """
        pass

    def render_prompt(self) -> None:
        """Render prompt templates with the current template data."""
        self.prompt = self.prompt_processor.render_prompts(self.prompt_template, self.template_data)

    # ---------------------------------
    # Model Execution
    # ---------------------------------

    def run_model(self) -> None:
        """Execute the model with the rendered prompt and configured parameters."""
        if self.agent_config.settings.system.debug.mode:
            self.result = self.agent_config.simulated_response
            return

        self._execute_model_generation()

    def _execute_model_generation(self) -> None:
        """Execute the actual model generation with configured parameters."""
        params = self._build_model_params()
        # `generate` may return str or raw bytes (for TTS).
        result = self.model.generate(self.prompt, **params)

        # Preserve raw bytes; strip only when we have text.
        if isinstance(result, str):
            result = result.strip()

        self.result = result

    def _build_model_params(self) -> Dict[str, Any]:
        """Build parameters for model generation."""
        params = self.agent_config.params.copy()
        params['agent_name'] = self.agent_name
        
        if self.images:
            params['images'] = self.images

        # Allow callers to supply raw audio for STT models via agent.run(audio=<bytes-or-path>)
        if 'audio' in self.template_data:
            params['audio'] = self.template_data['audio']

        return params

    # ---------------------------------
    # Result Handling
    # ---------------------------------

    def parse_result(self) -> None:
        """Handle text vs. binary results.

        • Text → optional structured parsing.
        • Bytes → save to disk via `AudioManager.save_tts_bytes` and return the file path.
        """

        # Audio bytes – treat as TTS output
        if isinstance(self.result, (bytes, bytearray)):
            saved_path = self.audio_manager.save_tts_bytes(self.result)
            self.parsed_result = saved_path  # downstream components will use the path
            return

        # Otherwise, default behaviour
        self.parsed_result = self.parsing_processor.parse_by_format(
            self.result, self.agent_config.parse_response_as
        )

    def post_process_result(self) -> None:
        """
        Extension point for additional processing after parsing the model's response.
        """
        pass

    def build_output(self) -> None:
        """
        Build the final output for the agent. By default, the output is set to the parsed result.
        Override this method in subclasses to implement custom output building.
        """
        self.output = self.parsed_result
