"""
Self-Modification Engine,  Allows the agent to inspect and modify its own code.

Features:
- Hot-reload modules at runtime
- Patch functions and methods safely
- Rollback failed modifications
- Audit trail of all changes
- Sandboxed code evaluation before applying
"""

import ast
import hashlib
import importlib
import inspect
import json
import logging
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class Modification:
    """Record of a single self-modification."""

    mod_id: str
    target_module: str
    target_name: str
    change_type: str  # "patch", "add", "remove", "replace"
    description: str
    old_source: Optional[str] = None
    new_source: Optional[str] = None
    applied_at: Optional[str] = None
    rolled_back: bool = False
    success: bool = False


@dataclass
class ModificationResult:
    """Result of a modification attempt."""

    success: bool
    modification: Modification
    error: Optional[str] = None


class SelfModificationEngine:
    """
    Enables the agent to safely modify its own behavior at runtime.

    Safety layers:
    1. AST validation,  new code must parse cleanly
    2. Signature check,  replacement functions must match the original signature
    3. Snapshot,  original code is saved before any change
    4. Rollback,  any failed modification is automatically reverted
    5. Audit log,  every change is recorded with timestamp and diff
    """

    def __init__(self, config=None):
        self.config = config
        self._history: List[Modification] = []
        self._snapshots: Dict[str, Any] = {}
        self._audit_file = Path.home() / ".sablecore" / "self_mod_audit.json"
        self._audit_file.parent.mkdir(parents=True, exist_ok=True)
        logger.info("🧬 Self-Modification Engine initialized")

    # ------------------------------------------------------------------
    # Inspection
    # ------------------------------------------------------------------

    def inspect_module(self, module_name: str) -> Dict[str, Any]:
        """Inspect a loaded module and return its structure."""
        try:
            mod = importlib.import_module(module_name)
            members = inspect.getmembers(mod)
            return {
                "module": module_name,
                "file": getattr(mod, "__file__", None),
                "classes": [name for name, obj in members if inspect.isclass(obj)],
                "functions": [name for name, obj in members if inspect.isfunction(obj)],
                "doc": inspect.getdoc(mod),
            }
        except Exception as e:
            return {"error": str(e)}

    def get_source(self, module_name: str, name: str) -> Optional[str]:
        """Get the source code of a function or class."""
        try:
            mod = importlib.import_module(module_name)
            obj = getattr(mod, name, None)
            if obj is None:
                return None
            return inspect.getsource(obj)
        except Exception as e:
            logger.error(f"Cannot get source for {module_name}.{name}: {e}")
            return None

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    def validate_code(self, code: str) -> Dict[str, Any]:
        """Validate Python code before applying it."""
        # Syntax check
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return {"valid": False, "error": f"SyntaxError: {e}"}

        # Extract defined names
        functions = [
            n.name for n in ast.walk(tree) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
        ]
        classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]

        return {
            "valid": True,
            "functions": functions,
            "classes": classes,
            "lines": len(code.splitlines()),
        }

    # ------------------------------------------------------------------
    # Modification
    # ------------------------------------------------------------------

    def patch_function(
        self,
        module_name: str,
        func_name: str,
        new_code: str,
        description: str = "",
    ) -> ModificationResult:
        """
        Replace a function in a loaded module with new code.

        Args:
            module_name: Dotted module path (e.g. 'opensable.core.agent')
            func_name: Name of the function to replace
            new_code: Python source of the replacement function
            description: Human-readable description of the change
        """
        mod_id = hashlib.sha256(
            f"{module_name}.{func_name}.{datetime.now().isoformat()}".encode()
        ).hexdigest()[:12]

        mod = Modification(
            mod_id=mod_id,
            target_module=module_name,
            target_name=func_name,
            change_type="patch",
            description=description,
        )

        try:
            # 1. Import module
            target_mod = importlib.import_module(module_name)
            original = getattr(target_mod, func_name, None)
            if original is None:
                mod.success = False
                return ModificationResult(
                    success=False, modification=mod, error=f"{func_name} not found in {module_name}"
                )

            # 2. Snapshot original
            mod.old_source = inspect.getsource(original)
            self._snapshots[mod_id] = original

            # 3. Validate new code
            validation = self.validate_code(new_code)
            if not validation["valid"]:
                return ModificationResult(
                    success=False, modification=mod, error=validation["error"]
                )

            # 4. Compile and extract the function
            namespace: Dict[str, Any] = {}
            exec(compile(new_code, f"<selfmod-{mod_id}>", "exec"), namespace)
            new_func = namespace.get(func_name)
            if new_func is None:
                return ModificationResult(
                    success=False, modification=mod, error=f"New code does not define '{func_name}'"
                )

            # 5. Apply
            setattr(target_mod, func_name, new_func)
            mod.new_source = new_code
            mod.applied_at = datetime.now().isoformat()
            mod.success = True

            self._history.append(mod)
            self._save_audit(mod)
            logger.info(f"🧬 Patched {module_name}.{func_name} (mod {mod_id})")

            return ModificationResult(success=True, modification=mod)

        except Exception as e:
            mod.success = False
            self._history.append(mod)
            return ModificationResult(success=False, modification=mod, error=str(e))

    def rollback(self, mod_id: str) -> bool:
        """Rollback a specific modification by its ID."""
        original = self._snapshots.get(mod_id)
        if original is None:
            logger.warning(f"No snapshot found for mod {mod_id}")
            return False

        # Find the modification record
        mod = next((m for m in self._history if m.mod_id == mod_id), None)
        if mod is None:
            return False

        try:
            target_mod = importlib.import_module(mod.target_module)
            setattr(target_mod, mod.target_name, original)
            mod.rolled_back = True
            logger.info(f"🔄 Rolled back {mod.target_module}.{mod.target_name}")
            return True
        except Exception as e:
            logger.error(f"Rollback failed: {e}")
            return False

    def hot_reload(self, module_name: str) -> bool:
        """Hot-reload an entire module from disk."""
        try:
            if module_name in sys.modules:
                mod = sys.modules[module_name]
                importlib.reload(mod)
                logger.info(f"🔃 Hot-reloaded: {module_name}")
                return True
            else:
                importlib.import_module(module_name)
                return True
        except Exception as e:
            logger.error(f"Hot-reload failed for {module_name}: {e}")
            return False

    # ------------------------------------------------------------------
    # Audit
    # ------------------------------------------------------------------

    def get_history(self) -> List[Modification]:
        """Return the full modification history."""
        return list(self._history)

    def _save_audit(self, mod: Modification):
        """Persist modification record to disk."""
        try:
            records = []
            if self._audit_file.exists():
                records = json.loads(self._audit_file.read_text())
            records.append(
                {
                    "mod_id": mod.mod_id,
                    "module": mod.target_module,
                    "name": mod.target_name,
                    "type": mod.change_type,
                    "description": mod.description,
                    "applied_at": mod.applied_at,
                    "success": mod.success,
                }
            )
            self._audit_file.write_text(json.dumps(records, indent=2))
        except Exception as e:
            logger.error(f"Failed to save audit log: {e}")
