"""Tool builder for generating new Python tools on-the-fly.

When the agent can't find an existing tool to solve a problem,
this module generates new tools using Claude.
"""

from __future__ import annotations

import ast
import re
from uuid import uuid4
from typing import Any

from promethian.core.types import Tool
from promethian.utils.claude import ClaudeClient


TOOL_BUILDER_SYSTEM = """You are a Python tool builder for a self-improving AI agent.

Your job is to create Python functions that solve specific problems. The tools you create:
1. Must be self-contained (all imports at the top)
2. Must have clear type hints
3. Must handle errors gracefully
4. Must return JSON-serializable results
5. Should be general enough to be reusable

When creating a tool:
1. Write clean, well-documented Python code
2. Include all necessary imports
3. Define the main function with clear parameters
4. Handle common error cases
5. Return useful results

IMPORTANT - Credential Handling:
If your tool uses any external APIs or services, you MUST:
1. Declare required environment variables in the required_env_vars field
2. Use standard naming conventions: OPENAI_API_KEY, STRIPE_SECRET_KEY, GITHUB_TOKEN, etc.
3. Access credentials via os.environ.get("VAR_NAME") in the code
4. NEVER hardcode API keys or secrets in the code

Output your tool as JSON with this structure:
{
    "name": "tool_name",
    "description": "When and how to use this tool",
    "function_name": "main_function_name",
    "code": "full Python code including imports",
    "parameters": {
        "type": "object",
        "properties": {
            "param1": {"type": "string", "description": "..."}
        },
        "required": ["param1"]
    },
    "dependencies": ["package1", "package2"],
    "required_env_vars": ["OPENAI_API_KEY"],
    "tags": ["tag1", "tag2"]
}"""


class ToolBuilderError(Exception):
    """Raised when tool building fails."""
    pass


class ToolBuilder:
    """
    Builds new Python tools using Claude.

    When the agent encounters a task that requires capabilities
    not available in existing tools, this builder creates new ones.
    """

    def __init__(self, claude: ClaudeClient):
        self.claude = claude

    async def build_tool(
        self,
        description: str,
        context: str | None = None,
        example_input: dict[str, Any] | None = None,
        example_output: Any | None = None,
    ) -> Tool:
        """
        Build a new tool from a description.

        Args:
            description: What the tool should do
            context: Additional context about the use case
            example_input: Example input parameters
            example_output: Expected output format

        Returns:
            A new Tool object ready for execution

        Raises:
            ToolBuilderError: If tool generation or validation fails
        """
        prompt = self._build_prompt(
            description=description,
            context=context,
            example_input=example_input,
            example_output=example_output,
        )

        try:
            tool_data = await self.claude.complete_json(
                prompt=prompt,
                system=TOOL_BUILDER_SYSTEM,
            )
        except Exception as e:
            raise ToolBuilderError(f"Failed to generate tool: {e}")

        # Validate the generated code
        code = tool_data.get("code", "")
        function_name = tool_data.get("function_name", "")

        validation_error = self._validate_code(code, function_name)
        if validation_error:
            raise ToolBuilderError(f"Generated code is invalid: {validation_error}")

        # Sanitize tool name to match Claude API requirements
        raw_name = tool_data.get("name", "unnamed_tool")
        sanitized_name = self._sanitize_tool_name(raw_name)

        # Build Tool object
        tool = Tool(
            id=f"tool_{uuid4().hex[:12]}",
            name=sanitized_name,
            description=tool_data.get("description", description),
            code=code,
            function_name=function_name,
            parameters=tool_data.get("parameters", {}),
            dependencies=tool_data.get("dependencies", []),
            required_env_vars=tool_data.get("required_env_vars", []),
            tags=tool_data.get("tags", []),
            source_task=description[:200],  # Track what created this tool
        )

        return tool

    async def improve_tool(
        self,
        tool: Tool,
        error: str,
        context: str | None = None,
    ) -> Tool | None:
        """
        Improve a tool that failed during execution.

        Args:
            tool: The tool that failed
            error: The error message/traceback
            context: Additional context about the failure

        Returns:
            An improved version of the tool, or None if no change needed
        """
        prompt = f"""Make MINIMAL fixes to this tool based on the error encountered.

IMPORTANT: Be extremely conservative. This tool has been proven to work in other cases.
Only fix the SPECIFIC error. Do NOT refactor, restructure, or "improve" unrelated code.
Change as few lines as possible.

Tool Name: {tool.name}
Description: {tool.description}
Function: {tool.function_name}

Current Code:
```python
{tool.code}
```

Error encountered:
{error}

{f"Context: {context}" if context else ""}

Rules for fixing:
1. PREFER returning no_change if the error was due to input/environment, not the code
2. Fix ONLY the specific line(s) that caused the error
3. Do NOT add extra error handling, logging, or "improvements"
4. Do NOT change the function signature, name, or parameters
5. Keep 95%+ of the code unchanged
6. If adding a fix, add the minimal code necessary

If the error was not due to a code bug (e.g., bad input, network issue, environment):
{{"no_change": true, "reason": "Error was due to X, not a code bug"}}

Otherwise respond with the minimally fixed tool JSON:
{{
    "name": "{tool.name}",
    "description": "{tool.description}",
    "function_name": "{tool.function_name}",
    "code": "...",
    "parameters": ...,
    "dependencies": [...],
    "tags": [...]
}}"""

        try:
            tool_data = await self.claude.complete_json(
                prompt=prompt,
                system=TOOL_BUILDER_SYSTEM,
            )
        except Exception as e:
            raise ToolBuilderError(f"Failed to improve tool: {e}")

        # Check if no change needed
        if tool_data.get("no_change"):
            return None

        # Validate the improved code
        code = tool_data.get("code", "")
        function_name = tool_data.get("function_name", tool.function_name)

        validation_error = self._validate_code(code, function_name)
        if validation_error:
            raise ToolBuilderError(f"Improved code is invalid: {validation_error}")

        # Build improved Tool object
        improved_tool = Tool(
            id=f"tool_{uuid4().hex[:12]}",
            name=tool_data.get("name", tool.name),
            description=tool_data.get("description", tool.description),
            code=code,
            function_name=function_name,
            parameters=tool_data.get("parameters", tool.parameters),
            dependencies=tool_data.get("dependencies", tool.dependencies),
            required_env_vars=tool_data.get("required_env_vars", tool.required_env_vars),
            tags=tool_data.get("tags", tool.tags),
            derived_from=tool.id,  # Track lineage
            source_task=tool.source_task,
            version=tool.version + 1,
        )

        return improved_tool

    async def build_from_example(
        self,
        example_code: str,
        task_description: str,
    ) -> Tool:
        """
        Build a tool by generalizing from example code.

        Useful when the agent has working code from an execution
        that should be turned into a reusable tool.

        Args:
            example_code: Working code to generalize
            task_description: What the code accomplishes

        Returns:
            A generalized Tool
        """
        prompt = f"""Turn this working code into a reusable tool.

Task: {task_description}

Working Code:
```python
{example_code}
```

Create a generalized, reusable tool that:
1. Extracts hardcoded values as parameters
2. Adds proper error handling
3. Has clear documentation
4. Can be used for similar tasks

Respond with the tool JSON."""

        try:
            tool_data = await self.claude.complete_json(
                prompt=prompt,
                system=TOOL_BUILDER_SYSTEM,
            )
        except Exception as e:
            raise ToolBuilderError(f"Failed to generalize code: {e}")

        code = tool_data.get("code", "")
        function_name = tool_data.get("function_name", "")

        validation_error = self._validate_code(code, function_name)
        if validation_error:
            raise ToolBuilderError(f"Generalized code is invalid: {validation_error}")

        return Tool(
            id=f"tool_{uuid4().hex[:12]}",
            name=tool_data.get("name", "generalized_tool"),
            description=tool_data.get("description", task_description),
            code=code,
            function_name=function_name,
            parameters=tool_data.get("parameters", {}),
            dependencies=tool_data.get("dependencies", []),
            required_env_vars=tool_data.get("required_env_vars", []),
            tags=tool_data.get("tags", []),
            source_task=task_description[:200],
        )

    def _build_prompt(
        self,
        description: str,
        context: str | None,
        example_input: dict[str, Any] | None,
        example_output: Any | None,
    ) -> str:
        """Build the prompt for tool generation."""
        parts = [f"Create a Python tool that: {description}"]

        if context:
            parts.append(f"\nContext: {context}")

        if example_input:
            parts.append(f"\nExample Input: {example_input}")

        if example_output:
            parts.append(f"\nExpected Output Format: {example_output}")

        parts.append("\nRespond with the complete tool JSON.")

        return "\n".join(parts)

    def _sanitize_tool_name(self, name: str) -> str:
        """Sanitize tool name to match Claude API requirements.

        Claude requires tool names to match: ^[a-zA-Z0-9_-]{1,64}$
        """
        # Replace spaces and invalid characters with underscores
        sanitized = re.sub(r'[^a-zA-Z0-9_-]', '_', name)
        # Remove consecutive underscores
        sanitized = re.sub(r'_+', '_', sanitized)
        # Remove leading/trailing underscores
        sanitized = sanitized.strip('_')
        # Truncate to 64 characters
        sanitized = sanitized[:64]
        # If empty, use a default
        if not sanitized:
            sanitized = "unnamed_tool"
        return sanitized

    def _validate_code(self, code: str, function_name: str) -> str | None:
        """
        Validate that the generated code is valid Python.

        Returns None if valid, error message if invalid.
        """
        if not code.strip():
            return "Code is empty"

        if not function_name:
            return "Function name is missing"

        # Check syntax
        try:
            ast.parse(code)
        except SyntaxError as e:
            return f"Syntax error: {e}"

        # Check that the function exists
        if not self._function_exists(code, function_name):
            return f"Function '{function_name}' not found in code"

        # Check for dangerous patterns (basic security)
        dangerous_patterns = [
            r'\bos\.system\b',
            r'\bsubprocess\.call\b.*shell\s*=\s*True',
            r'\beval\b\s*\(',
            r'\bexec\b\s*\(',
            r'__import__\s*\(',
        ]

        for pattern in dangerous_patterns:
            if re.search(pattern, code):
                return f"Potentially dangerous pattern detected: {pattern}"

        return None

    def _function_exists(self, code: str, function_name: str) -> bool:
        """Check if a function with the given name exists in the code."""
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef) and node.name == function_name:
                    return True
                if isinstance(node, ast.AsyncFunctionDef) and node.name == function_name:
                    return True
            return False
        except SyntaxError:
            return False
