from typing import Dict, Any
from loguru import logger
from utils.llm_orchestrator import LLMOrchestrator


class CodeGenerationAgent:
    def __init__(self, llm_api_key: str):
        """Initialize the Code Generation Agent."""
        logger.info("Initializing CodeGenerationAgent")
        self.llm_orchestrator = LLMOrchestrator(llm_api_key)
        self.capabilities = [
            "code_generation",
            "code_completion",
            "code_modification",
            "code_refactoring"
        ]
        self.setup_logger()

    def setup_logger(self):
        """Configure logging for the agent."""
        logger.add("logs/code_generation_agent.log", rotation="500 MB")

    async def generate_code(self, description: str,
                            language: str = "python") -> Dict[str, Any]:
        """Generate code based on a natural language description."""
        logger.info(f"Generating code for description: {description}")
        try:
            prompt = f"""
            Generate {language} code that accomplishes the following:
            {description}

            Provide only the code, without any introductory or explanatory text.
            """
            code = await self.llm_orchestrator.generate_completion(prompt)

            logger.info(f"Code generated successfully:\n{code}")
            return {
                "status": "success",
                "code": code
            }

        except Exception as e:
            logger.error(f"Error generating code: {str(e)}")
            return {
                "status": "error",
                "message": str(e)
            }