"""
PyGEAI-based agent implementation.

This module provides a concrete agent implementation using the PyGEAI SDK,
enabling AI-powered task execution with GeneXus Enterprise AI.
"""

import logging
from typing import Any, Dict, Optional

from pygeai_orchestration.core.base.agent import BaseAgent, AgentConfig
from pygeai.core.base.session import Session
from pygeai.chat.managers import ChatManager
from pygeai.core.models import LlmSettings, ChatMessageList, ChatMessage
from pygeai.core.common.config import get_settings

logger = logging.getLogger("pygeai_orchestration")


class GEAIAgent(BaseAgent):
    """
    Agent implementation using PyGEAI SDK.

    This agent leverages the PyGEAI SDK to execute tasks using GeneXus
    Enterprise AI models, providing AI-powered reasoning and generation.

    The agent:
    - Integrates with PyGEAI Session and ChatManager
    - Supports custom system prompts and model configuration
    - Maintains execution history
    - Automatically loads credentials from ~/.geai/credentials
    """

    def __init__(self, config: AgentConfig, session=None, alias: Optional[str] = None):
        """
        Initialize the GEAI agent with PyGEAI session.

        :param config: AgentConfig - Agent configuration.
        :param session: Optional PyGEAI Session - Existing session to use.
        :param alias: Optional[str] - Credentials alias (default: 'default').
        """
        super().__init__(config)
        
        if session:
            self._session = session
        else:
            alias = alias or 'default'
            settings = get_settings()
            api_key = settings.get_api_key(alias)
            base_url = settings.get_base_url(alias)
            self._session = Session(api_key=api_key, base_url=base_url)
        
        self._chat_manager = ChatManager(self._session)

    async def execute(self, task: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        Execute a task using PyGEAI generation.

        :param task: str - Task description or instruction.
        :param context: Optional[Dict[str, Any]] - Execution context.
        :return: Dict[str, Any] - Execution result with success status and output.
        """
        logger.debug(f"GEAIAgent '{self.name}' executing task")

        try:
            result = await self.generate(task, context=context)

            execution_result = {
                "success": True,
                "result": result,
                "agent": self.name,
                "metadata": context or {},
            }

            self.add_to_history(execution_result)
            return execution_result

        except Exception as e:
            logger.error(f"GEAIAgent execution failed: {str(e)}")
            error_result = {"success": False, "result": None, "error": str(e), "agent": self.name}
            self.add_to_history(error_result)
            return error_result

    async def generate(self, prompt: str, **kwargs) -> str:
        """
        Generate a response using PyGEAI chat completion.

        :param prompt: str - Input prompt for generation.
        :param kwargs: Additional generation parameters (messages, etc.).
        :return: str - Generated text response.
        :raises Exception: If generation fails.
        """
        try:
            messages = kwargs.get("messages", [])
            if not messages:
                messages = [{"role": "user", "content": prompt}]

            if self.config.system_prompt:
                messages.insert(0, {"role": "system", "content": self.config.system_prompt})

            chat_messages = ChatMessageList(
                messages=[ChatMessage(role=msg["role"], content=msg["content"]) for msg in messages]
            )

            llm_settings = LlmSettings(
                temperature=self.config.temperature, max_tokens=self.config.max_tokens
            )

            response = self._chat_manager.chat_completion(
                model=self.config.model, messages=chat_messages, llm_settings=llm_settings
            )

            response_text = response.choices[0].message.content

            self.add_to_history({"type": "generation", "prompt": prompt, "response": response_text})

            return response_text

        except Exception as e:
            logger.error(f"Generation failed: {str(e)}")
            raise
