"""
Code modification module for automatic strategy code generation and modification
"""

import os
import ast
from pathlib import Path
from dataclasses import dataclass, field
from typing import Dict, Any, List, Optional, Tuple
from anthropic import Anthropic

from ..models.decision import CodeModification, ImprovementPlan


@dataclass
class FileModification:
    """
    File modification specification

    Attributes:
        file_path: Path to file to be modified
        modification_type: Type of modification ('update' | 'create' | 'delete')
        original_content: Original file content (if exists)
        new_content: New file content
        line_range: Optional line range for partial modifications
        description: Modification description
    """

    file_path: str
    modification_type: str  # 'update' | 'create' | 'delete'
    original_content: Optional[str]
    new_content: str
    line_range: Optional[Tuple[int, int]] = None
    description: str = ""


@dataclass
class CodeContext:
    """
    Code context for understanding file structure

    Attributes:
        file_path: Path to the file
        file_content: Full file content
        imports: List of import statements
        classes: List of class names
        functions: List of function names
        dependencies: List of dependency file paths
    """

    file_path: str
    file_content: str
    imports: List[str] = field(default_factory=list)
    classes: List[str] = field(default_factory=list)
    functions: List[str] = field(default_factory=list)
    dependencies: List[str] = field(default_factory=list)


@dataclass
class ModificationResult:
    """
    Result of code modification operation

    Attributes:
        success: Whether modification was successful
        modified_files: List of modified file paths
        errors: List of error messages
        warnings: List of warning messages
        git_commit: Git commit SHA if committed
        timestamp: Modification timestamp
    """

    success: bool
    modified_files: List[str] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    git_commit: Optional[str] = None
    timestamp: int = 0


class CodeModifier:
    """
    Code modifier that uses Claude to generate and modify strategy code

    This class integrates with Claude API and Git to automatically modify
    strategy code based on improvement plans.
    """

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

        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
        """
        self.client = claude_client
        self.workspace = Path(workspace_path)
        self.git = git_manager
        self.model = model_version

    def parse_improvement_plan(self, plan: ImprovementPlan) -> List[FileModification]:
        """
        Parse improvement plan into file modifications

        Args:
            plan: Improvement plan from decision engine

        Returns:
            List of file modifications to apply
        """
        file_modifications = []

        for mod in plan.modifications:
            # Resolve file path
            file_path = self._resolve_file_path(mod.file_path)

            # Read current content if file exists
            original_content = None
            if file_path.exists():
                with open(file_path, "r", encoding="utf-8") as f:
                    original_content = f.read()

            # Determine modification type
            if not file_path.exists():
                mod_type = "create"
            elif mod.modification_type == "delete":
                mod_type = "delete"
            else:
                mod_type = "update"

            file_modifications.append(
                FileModification(
                    file_path=str(file_path),
                    modification_type=mod_type,
                    original_content=original_content,
                    new_content="",  # Will be generated
                    description=mod.description,
                )
            )

        return file_modifications

    def locate_code_files(self, modifications: List[CodeModification]) -> Dict[str, str]:
        """
        Locate code files that need to be modified

        Args:
            modifications: List of code modifications

        Returns:
            Dictionary mapping modification target to file path
        """
        file_mapping = {}

        for mod in modifications:
            # Try to resolve the file path
            file_path = self._resolve_file_path(mod.file_path)

            if file_path.exists():
                file_mapping[mod.target_component] = str(file_path)
            else:
                # Try to find the file by searching for the component
                found_path = self._search_for_component(mod.target_component)
                if found_path:
                    file_mapping[mod.target_component] = found_path

        return file_mapping

    def generate_code(self, modification: CodeModification, context: CodeContext) -> str:
        """
        Generate modified code using Claude

        Args:
            modification: Code modification specification
            context: Code context for understanding the file

        Returns:
            Generated code content
        """
        from .code_prompts import CodeGenerationTemplates

        # Build prompt for code generation
        prompt = CodeGenerationTemplates.GENERATE_CODE.format(
            file_path=context.file_path,
            current_code=context.file_content,
            modification_description=modification.description,
            modification_type=modification.modification_type,
            target_component=modification.target_component,
            current_value=modification.current_value,
            new_value=modification.new_value,
        )

        # Call Claude API
        response = self.client.messages.create(
            model=self.model,
            max_tokens=4096,
            messages=[{"role": "user", "content": prompt}],
        )

        # Extract code from response
        generated_code = response.content[0].text

        # Clean up code blocks if present
        if "```python" in generated_code:
            # Extract code from markdown code block
            start = generated_code.find("```python") + 9
            end = generated_code.rfind("```")
            generated_code = generated_code[start:end].strip()

        return generated_code

    def apply_modifications(self, file_modifications: List[FileModification]) -> ModificationResult:
        """
        Apply code modifications to files

        Args:
            file_modifications: List of file modifications

        Returns:
            Modification result with success status and details
        """
        result = ModificationResult(success=True, timestamp=int(os.times().elapsed))
        modified_files = []
        errors = []
        warnings = []

        try:
            for file_mod in file_modifications:
                file_path = Path(file_mod.file_path)

                if file_mod.modification_type == "delete":
                    # Delete file
                    if file_path.exists():
                        file_path.unlink()
                        modified_files.append(str(file_path))
                    else:
                        warnings.append(f"File not found for deletion: {file_path}")

                elif file_mod.modification_type == "create":
                    # Create new file
                    file_path.parent.mkdir(parents=True, exist_ok=True)
                    with open(file_path, "w", encoding="utf-8") as f:
                        f.write(file_mod.new_content)
                    modified_files.append(str(file_path))

                elif file_mod.modification_type == "update":
                    # Update existing file
                    if not file_path.exists():
                        errors.append(f"File not found for update: {file_path}")
                        continue

                    with open(file_path, "w", encoding="utf-8") as f:
                        f.write(file_mod.new_content)
                    modified_files.append(str(file_path))

            result.modified_files = modified_files
            result.warnings = warnings

            # Commit changes if successful
            if modified_files and not errors:
                commit_sha = self.git.commit_changes(
                    files=modified_files, message="Auto-generated code modifications"
                )
                result.git_commit = commit_sha

        except Exception as e:
            result.success = False
            errors.append(f"Error applying modifications: {str(e)}")

        result.errors = errors
        if errors:
            result.success = False

        return result

    def _resolve_file_path(self, file_path: str) -> Path:
        """
        Resolve file path relative to workspace

        Args:
            file_path: File path (relative or absolute)

        Returns:
            Resolved Path object
        """
        path = Path(file_path)
        if path.is_absolute():
            return path
        return self.workspace / path

    def _search_for_component(self, component_name: str) -> Optional[str]:
        """
        Search for a component (class/function) in the workspace

        Args:
            component_name: Name of component to find

        Returns:
            File path if found, None otherwise
        """
        # Search Python files in workspace
        for py_file in self.workspace.rglob("*.py"):
            try:
                with open(py_file, "r", encoding="utf-8") as f:
                    content = f.read()

                # Parse AST to find classes and functions
                tree = ast.parse(content)
                for node in ast.walk(tree):
                    if isinstance(node, (ast.ClassDef, ast.FunctionDef)):
                        if node.name == component_name:
                            return str(py_file)
            except Exception:
                # Skip files that can't be parsed
                continue

        return None

    def extract_code_context(self, file_path: str) -> CodeContext:
        """
        Extract code context from a file

        Args:
            file_path: Path to the file

        Returns:
            CodeContext with file structure information
        """
        path = Path(file_path)
        if not path.exists():
            return CodeContext(file_path=file_path, file_content="")

        with open(path, "r", encoding="utf-8") as f:
            content = f.read()

        context = CodeContext(file_path=file_path, file_content=content)

        try:
            # Parse AST
            tree = ast.parse(content)

            # Extract imports
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for alias in node.names:
                        context.imports.append(alias.name)
                elif isinstance(node, ast.ImportFrom):
                    if node.module:
                        context.imports.append(node.module)

            # Extract classes and functions
            for node in ast.iter_child_nodes(tree):
                if isinstance(node, ast.ClassDef):
                    context.classes.append(node.name)
                elif isinstance(node, ast.FunctionDef):
                    context.functions.append(node.name)

        except SyntaxError:
            # If file has syntax errors, return partial context
            pass

        return context
