"""
Code Modification Tool Wrapper - Tool interface for Agent to access code modification functionality

This module wraps the CodeModifier as a BaseTool for Agent integration.
"""

import logging
from typing import Dict, Any, Optional

from openfund_agent.tools.base import BaseTool, register_tool
from openfund_agent.modification.modifier import CodeModifier, FileModification, CodeContext
from openfund_agent.models.decision import ImprovementPlan, CodeModification as CodeMod


@register_tool
class CodeModificationTool(BaseTool):
    """
    Code modification tool wrapper for Agent

    This tool provides Agent with access to code generation and modification
    functionality through a standardized Tool interface.
    """

    def __init__(
        self,
        claude_client,
        workspace_path: str,
        git_manager,
        model_version: str = "claude-3-5-sonnet-20241022",
    ):
        """
        Initialize code modification tool

        Args:
            claude_client: Anthropic Claude API client
            workspace_path: Path to workspace root
            git_manager: Git manager instance
            model_version: Claude model version to use
        """
        super().__init__(
            name="code_modification",
            description="Generate and modify strategy code based on improvement plans",
        )

        self.logger = logging.getLogger(__name__)
        self.modifier = CodeModifier(
            claude_client=claude_client,
            workspace_path=workspace_path,
            git_manager=git_manager,
            model_version=model_version,
        )

        self.logger.info("CodeModificationTool initialized")

    def _define_parameters(self) -> Dict[str, Any]:
        """
        Define parameter schema for code modification tool

        Returns:
            Parameter schema dictionary
        """
        return {
            "action": {
                "type": str,
                "required": True,
                "enum": [
                    "parse_plan",
                    "locate_files",
                    "generate_code",
                    "apply_modifications",
                    "extract_context",
                ],
                "description": "Action to perform: parse_plan, locate_files, generate_code, apply_modifications, or extract_context",
            },
            "improvement_plan": {
                "type": dict,
                "required": False,
                "description": "Improvement plan dictionary (required for parse_plan)",
            },
            "modifications": {
                "type": list,
                "required": False,
                "description": "List of code modification dictionaries (required for locate_files)",
            },
            "modification": {
                "type": dict,
                "required": False,
                "description": "Single code modification dictionary (required for generate_code)",
            },
            "context": {
                "type": dict,
                "required": False,
                "description": "Code context dictionary (required for generate_code)",
            },
            "file_modifications": {
                "type": list,
                "required": False,
                "description": "List of file modification dictionaries (required for apply_modifications)",
            },
            "file_path": {
                "type": str,
                "required": False,
                "description": "File path (required for extract_context)",
            },
        }

    def execute(self, parameters: Dict[str, Any]) -> Any:
        """
        Execute code modification tool

        Args:
            parameters: Tool parameters
                - action: Action to perform
                - improvement_plan: Improvement plan (for parse_plan)
                - modifications: Code modifications (for locate_files)
                - modification: Single modification (for generate_code)
                - context: Code context (for generate_code)
                - file_modifications: File modifications (for apply_modifications)
                - file_path: File path (for extract_context)

        Returns:
            Execution result based on action type

        Raises:
            ValueError: If parameters are invalid
        """
        action = parameters["action"]

        self.logger.info(f"Executing code modification tool: action={action}")

        try:
            if action == "parse_plan":
                return self._parse_plan(parameters)
            elif action == "locate_files":
                return self._locate_files(parameters)
            elif action == "generate_code":
                return self._generate_code(parameters)
            elif action == "apply_modifications":
                return self._apply_modifications(parameters)
            elif action == "extract_context":
                return self._extract_context(parameters)
            else:
                raise ValueError(f"Invalid action: '{action}'")

        except Exception as e:
            self.logger.error(f"Code modification tool execution failed: {str(e)}")
            raise

    def _parse_plan(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        Parse improvement plan into file modifications

        Args:
            parameters: Parameters including improvement_plan

        Returns:
            Dictionary with file modifications
        """
        plan_dict = parameters.get("improvement_plan")
        if not plan_dict:
            raise ValueError("improvement_plan is required for parse_plan action")

        # Convert dict to ImprovementPlan object
        plan = self._dict_to_improvement_plan(plan_dict)

        # Parse plan
        file_modifications = self.modifier.parse_improvement_plan(plan)

        # Convert to serializable format
        return {
            "success": True,
            "file_modifications": [
                {
                    "file_path": fm.file_path,
                    "modification_type": fm.modification_type,
                    "original_content": fm.original_content,
                    "new_content": fm.new_content,
                    "line_range": fm.line_range,
                    "description": fm.description,
                }
                for fm in file_modifications
            ],
            "count": len(file_modifications),
        }

    def _locate_files(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        Locate code files that need to be modified

        Args:
            parameters: Parameters including modifications list

        Returns:
            Dictionary with file mapping
        """
        modifications_list = parameters.get("modifications")
        if not modifications_list:
            raise ValueError("modifications is required for locate_files action")

        # Convert dicts to CodeModification objects
        modifications = [self._dict_to_code_modification(m) for m in modifications_list]

        # Locate files
        file_mapping = self.modifier.locate_code_files(modifications)

        return {
            "success": True,
            "file_mapping": file_mapping,
            "count": len(file_mapping),
        }

    def _generate_code(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        Generate modified code using Claude

        Args:
            parameters: Parameters including modification and context

        Returns:
            Dictionary with generated code
        """
        modification_dict = parameters.get("modification")
        context_dict = parameters.get("context")

        if modification_dict is None:
            raise ValueError("modification is required for generate_code action")
        if context_dict is None:
            raise ValueError("context is required for generate_code action")

        # Convert dicts to objects
        modification = self._dict_to_code_modification(modification_dict)
        context = self._dict_to_code_context(context_dict)

        # Generate code
        generated_code = self.modifier.generate_code(modification, context)

        return {
            "success": True,
            "generated_code": generated_code,
            "file_path": context.file_path,
        }

    def _apply_modifications(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        Apply code modifications to files

        Args:
            parameters: Parameters including file_modifications list

        Returns:
            Dictionary with modification result
        """
        file_modifications_list = parameters.get("file_modifications")
        if not file_modifications_list:
            raise ValueError("file_modifications is required for apply_modifications action")

        # Convert dicts to FileModification objects
        file_modifications = [self._dict_to_file_modification(fm) for fm in file_modifications_list]

        # Apply modifications
        result = self.modifier.apply_modifications(file_modifications)

        return {
            "success": result.success,
            "modified_files": result.modified_files,
            "errors": result.errors,
            "warnings": result.warnings,
            "git_commit": result.git_commit,
            "timestamp": result.timestamp,
        }

    def _extract_context(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        Extract code context from a file

        Args:
            parameters: Parameters including file_path

        Returns:
            Dictionary with code context
        """
        file_path = parameters.get("file_path")
        if not file_path:
            raise ValueError("file_path is required for extract_context action")

        # Extract context
        context = self.modifier.extract_code_context(file_path)

        return {
            "success": True,
            "context": {
                "file_path": context.file_path,
                "file_content": context.file_content,
                "imports": context.imports,
                "classes": context.classes,
                "functions": context.functions,
                "dependencies": context.dependencies,
            },
        }

    def _dict_to_improvement_plan(self, plan_dict: Dict[str, Any]) -> ImprovementPlan:
        """
        Convert dictionary to ImprovementPlan object

        Args:
            plan_dict: Plan dictionary

        Returns:
            ImprovementPlan object
        """
        modifications = [
            self._dict_to_code_modification(m) for m in plan_dict.get("modifications", [])
        ]

        return ImprovementPlan(
            plan_id=plan_dict.get("plan_id", ""),
            issues_addressed=plan_dict.get("issues_addressed", []),
            modifications=modifications,
            expected_impact=plan_dict.get("expected_impact", {}),
            risk_level=plan_dict.get("risk_level", "low"),
            requires_review=plan_dict.get("requires_review", False),
            rationale=plan_dict.get("rationale", ""),
            timestamp=plan_dict.get("timestamp", 0),
        )

    def _dict_to_code_modification(self, mod_dict: Dict[str, Any]) -> CodeMod:
        """
        Convert dictionary to CodeModification object

        Args:
            mod_dict: Modification dictionary

        Returns:
            CodeModification object
        """
        return CodeMod(
            file_path=mod_dict.get("file_path", ""),
            modification_type=mod_dict.get("modification_type", "parameter"),
            target_component=mod_dict.get("target_component", ""),
            current_value=mod_dict.get("current_value"),
            new_value=mod_dict.get("new_value"),
            description=mod_dict.get("description", ""),
        )

    def _dict_to_code_context(self, context_dict: Dict[str, Any]) -> CodeContext:
        """
        Convert dictionary to CodeContext object

        Args:
            context_dict: Context dictionary

        Returns:
            CodeContext object
        """
        return CodeContext(
            file_path=context_dict.get("file_path", ""),
            file_content=context_dict.get("file_content", ""),
            imports=context_dict.get("imports", []),
            classes=context_dict.get("classes", []),
            functions=context_dict.get("functions", []),
            dependencies=context_dict.get("dependencies", []),
        )

    def _dict_to_file_modification(self, fm_dict: Dict[str, Any]) -> FileModification:
        """
        Convert dictionary to FileModification object

        Args:
            fm_dict: File modification dictionary

        Returns:
            FileModification object
        """
        return FileModification(
            file_path=fm_dict.get("file_path", ""),
            modification_type=fm_dict.get("modification_type", "update"),
            original_content=fm_dict.get("original_content"),
            new_content=fm_dict.get("new_content", ""),
            line_range=fm_dict.get("line_range"),
            description=fm_dict.get("description", ""),
        )
