"""
AI Agent module for shell command generation
"""

from textwrap import dedent
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel

from .config import get_model, setup_environment, generate_modern_commands_prompt
from .models import Answer

def create_system_prompt() -> str:
    """Create system prompt with modern commands preferences"""
    base_prompt = dedent(
        """\
        You are a professional developer specializing in shell commands.
        Your task is to generate the correct shell commands based on the
        user's request.

        IMPORTANT: ALWAYS USE THE SAME LANGUAGE AS THE USER PROMPT IN
        YOUR RESPONSE.

        Process:

        1. Think Aloud: Use the `think` function to explain your reasoning.
        Justify why you chose a particular command, considering efficiency,
        safety, and best practices.

        2. Provide the Final Command: Use the `answer` function to present
        the final shell command concisely.
        """
    )

    # Add modern commands preferences
    modern_commands_prompt = generate_modern_commands_prompt()

    return base_prompt + modern_commands_prompt

def create_agent() -> Agent:
    """Create and configure the AI agent"""
    # Setup environment variables
    setup_environment()

    # Create OpenAI compatible model
    model = OpenAIModel(get_model())

    # Create agent with dynamic system prompt
    agent = Agent(
        model=model,
        system_prompt=create_system_prompt(),
        output_type=Answer,
    )
    
    # Register tools
    @agent.tool_plain
    def think(s: str) -> None:
        """Communicate your thought process to the user.

        Args:
            s (str): A description of your reasoning or decision-making process.
        """
        print(f"(AI Thinking): {s}\n")

    @agent.tool_plain
    def answer(success: bool, cmd: str | None, failure: str | None) -> Answer:
        """Provide the final shell command or explain why it couldn't be generated.

        Args:
            success (bool): Indicates whether a shell command was successfully generated.
            cmd (str | None): The generated shell command if `success` is True.
                It must be a single-line command. If `success` is False, this should be None.
            failure (str | None): If `success` is False, provide a reason why the command
                could not be generated. If `success` is True, this should be None.

        Returns:
            Answer: A structured response that will be processed for the user.
        """
        return Answer(success, cmd, failure)
    
    return agent