"""
Patch Agent - Apply constrained code edits with pre/post test gates

WO-ECHO-PATCH-AGENT-01: Agent for patching files with safety guarantees.

Features:
- File allowlist enforcement (only helix_cdc/, api/, tools/, tests/)
- Pre-test gate: baseline tests must pass before patching
- Post-test gate: tests must pass after patching or rollback
- Backup/rollback: automatic restore on post-test failure
- Receipt generation with chain-hash linking pre/post test receipts

Usage:
    from helix_cdc.echo.agents.patch_agent import PatchAgent
    from helix_cdc.echo.agents import register_agent

    register_agent(PatchAgent())
"""

import asyncio
import difflib
import hashlib
import json
import shutil
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from . import (
    AgentResult,
    DispatchableAgent,
    PROJECT_ROOT,
    generate_agent_id,
)

# Import path validation and intent detection from ops helpers
from helix_cdc.echo.ops.helpers import (
    validate_patch_path,
    detect_patch_intent,
    generate_receipt_id,
    ALLOWED_PATCH_PATHS,
)


class PatchAgent(DispatchableAgent):
    """
    Agent for applying code patches with test gates.

    Enforces safety through:
    - File allowlist (only modify allowed paths)
    - Pre-test gate (baseline must pass)
    - Post-test gate (rollback if fails)
    - Receipt chain linking test receipts

    Safety:
    - No direct subprocess calls; uses TestAgent for execution
    - All file paths validated against PROJECT_ROOT
    - Automatic backup before modification
    - Diff size limits to prevent large changes
    """

    # Class-level configuration
    MAX_DIFF_LINES = 500  # Safety limit on patch size
    DEFAULT_TEST_COMMAND = "pytest tests/ -q"

    @property
    def agent_type(self) -> str:
        return "patch"

    def __init__(
        self,
        receipt_dir: str = "receipts/patch_agent",
        backup_dir: str = "receipts/patch_agent/.backups",
        timeout_seconds: int = 120,
    ):
        """
        Initialize the patch agent.

        Args:
            receipt_dir: Directory for receipt storage (relative to PROJECT_ROOT)
            backup_dir: Directory for file backups (relative to PROJECT_ROOT)
            timeout_seconds: Maximum time for test execution
        """
        self.receipt_dir = PROJECT_ROOT / receipt_dir
        self.backup_dir = PROJECT_ROOT / backup_dir
        self.timeout_seconds = timeout_seconds
        self._test_agent = None

    def available(self) -> bool:
        """Check if agent is available (PROJECT_ROOT writable)."""
        try:
            test_file = PROJECT_ROOT / ".patch_agent_check"
            test_file.touch()
            test_file.unlink()
            return True
        except Exception:
            return False

    def _get_test_agent(self):
        """Get or create TestAgent instance."""
        if self._test_agent is None:
            from .test_agent import TestAgent
            self._test_agent = TestAgent(timeout_seconds=self.timeout_seconds)
        return self._test_agent

    def _parse_patch_intent(
        self, query: str, context: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Parse patch intent from query and context.

        Returns dict with:
            - files: List of file paths to modify
            - operation: Type of operation (add, remove, replace, etc.)
            - content: Content to add/modify (if applicable)
            - patch_spec: Full patch specification from context
        """
        # Check if patch_spec is provided in context
        patch_spec = context.get("patch_spec", {})

        # Try to detect intent from prompt if not provided
        if not patch_spec:
            detection = detect_patch_intent(query)
            files = detection.get("detected_files", [])
            patch_spec = {
                "files": files,
                "operation": "unknown",
                "query": query,
            }
        else:
            # Normalize patch_spec
            if "file" in patch_spec and "files" not in patch_spec:
                patch_spec["files"] = [patch_spec["file"]]
            if "files" not in patch_spec:
                patch_spec["files"] = []

        return patch_spec

    def _validate_file_path(self, file_path: str) -> Tuple[bool, str]:
        """Validate file path against allowlist."""
        return validate_patch_path(file_path)

    async def _run_test_gate(self, phase: str) -> AgentResult:
        """
        Run test gate via TestAgent.

        Args:
            phase: "pre" or "post" for logging

        Returns:
            AgentResult from TestAgent
        """
        test_agent = self._get_test_agent()
        query = f"run tests for {phase}-patch verification"
        context = {
            "parsed_command": {
                "command": "pytest",
                "args": ["tests/", "-q"],
            }
        }
        return await test_agent.execute(query, context)

    def _create_backups(self, files: List[str]) -> Dict[str, str]:
        """
        Create backups of files before modification.

        Args:
            files: List of file paths (relative to PROJECT_ROOT)

        Returns:
            Dict mapping original path to backup path
        """
        self.backup_dir.mkdir(parents=True, exist_ok=True)

        backups = {}
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")

        for file_path in files:
            source = PROJECT_ROOT / file_path
            if source.exists():
                # Generate backup filename
                safe_name = file_path.replace("/", "_").replace("\\", "_")
                backup_name = f"{safe_name}_{timestamp}.bak"
                backup_path = self.backup_dir / backup_name

                shutil.copy2(source, backup_path)
                backups[file_path] = str(backup_path.relative_to(PROJECT_ROOT))

        return backups

    def _restore_backups(self, backups: Dict[str, str]) -> List[str]:
        """
        Restore files from backups.

        Args:
            backups: Dict mapping original path to backup path

        Returns:
            List of restored file paths
        """
        restored = []
        for original_path, backup_path in backups.items():
            source = PROJECT_ROOT / backup_path
            dest = PROJECT_ROOT / original_path
            if source.exists():
                shutil.copy2(source, dest)
                restored.append(original_path)
        return restored

    def _apply_patch(self, patch_spec: Dict[str, Any]) -> Dict[str, Any]:
        """
        Apply the patch to files.

        Args:
            patch_spec: Patch specification with files, operation, content

        Returns:
            Dict with diffs, lines_added, lines_removed, files_modified
        """
        operation = patch_spec.get("operation", "unknown")
        files = patch_spec.get("files", [])
        content = patch_spec.get("content", "")

        diffs = {}
        total_added = 0
        total_removed = 0
        files_modified = []

        for file_path in files:
            full_path = PROJECT_ROOT / file_path

            # Read original content
            original_content = ""
            if full_path.exists():
                original_content = full_path.read_text(errors="ignore")

            # Apply operation
            new_content = original_content
            if operation == "add_comment":
                # Add a comment at the top
                new_content = f"# {content}\n" + original_content
            elif operation == "append":
                new_content = original_content + "\n" + content
            elif operation == "prepend":
                new_content = content + "\n" + original_content
            elif operation == "replace":
                old_text = patch_spec.get("old_text", "")
                if old_text and old_text in original_content:
                    new_content = original_content.replace(old_text, content, 1)
            elif operation == "replace_all":
                old_text = patch_spec.get("old_text", "")
                if old_text:
                    new_content = original_content.replace(old_text, content)
            else:
                # Default: treat content as full file replacement if provided
                if content:
                    new_content = content

            # Compute diff
            original_lines = original_content.splitlines(keepends=True)
            new_lines = new_content.splitlines(keepends=True)

            diff = list(difflib.unified_diff(
                original_lines,
                new_lines,
                fromfile=f"a/{file_path}",
                tofile=f"b/{file_path}",
            ))
            diff_text = "".join(diff)

            # Count changes
            added = sum(1 for line in diff if line.startswith("+") and not line.startswith("+++"))
            removed = sum(1 for line in diff if line.startswith("-") and not line.startswith("---"))

            # Check diff size limit
            if added + removed > self.MAX_DIFF_LINES:
                raise ValueError(
                    f"Patch too large: {added + removed} lines (max {self.MAX_DIFF_LINES})"
                )

            # Write new content
            if new_content != original_content:
                full_path.parent.mkdir(parents=True, exist_ok=True)
                full_path.write_text(new_content)
                files_modified.append(file_path)
                diffs[file_path] = diff_text
                total_added += added
                total_removed += removed

        return {
            "diffs": diffs,
            "lines_added": total_added,
            "lines_removed": total_removed,
            "files_modified": files_modified,
        }

    def _compute_file_hash(self, file_path: str) -> str:
        """Compute SHA256 hash of file content."""
        full_path = PROJECT_ROOT / file_path
        if full_path.exists():
            content = full_path.read_bytes()
            return hashlib.sha256(content).hexdigest()
        return ""

    def _compute_chain_hash(
        self,
        prompt_sha: str,
        diff_sha: str,
        pre_hash: str,
        post_hash: str,
        status: str,
    ) -> str:
        """
        Compute chain hash linking all verification hashes.

        chain_hash = sha256(prompt_sha || diff_sha || pre_hash || post_hash || status)
        """
        data = f"{prompt_sha}|{diff_sha}|{pre_hash}|{post_hash}|{status}"
        return hashlib.sha256(data.encode()).hexdigest()

    def _build_receipt(
        self,
        query: str,
        patch_spec: Dict[str, Any],
        pre_receipt: Optional[str],
        post_receipt: Optional[str],
        file_diffs: Dict[str, str],
        file_hashes_before: Dict[str, str],
        file_hashes_after: Dict[str, str],
        backup_manifest: Dict[str, str],
        lines_added: int,
        lines_removed: int,
        status: str,
        error: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Build a receipt for the patch execution.

        Returns receipt dict with chain-hash linking pre/post test receipts.
        """
        receipt_id = generate_receipt_id()
        timestamp = datetime.now(timezone.utc)

        # Compute hashes
        prompt_sha = hashlib.sha256(query.encode()).hexdigest()

        # Diff hash (hash of all diffs combined)
        all_diffs = "\n".join(file_diffs.values())
        diff_sha = hashlib.sha256(all_diffs.encode()).hexdigest()

        # Get pre/post chain hashes from receipts
        pre_chain_hash = ""
        post_chain_hash = ""
        if pre_receipt:
            pre_receipt_path = PROJECT_ROOT / pre_receipt
            if pre_receipt_path.exists():
                with open(pre_receipt_path) as f:
                    pre_data = json.load(f)
                    pre_chain_hash = pre_data.get("chain_hash", "")

        if post_receipt:
            post_receipt_path = PROJECT_ROOT / post_receipt
            if post_receipt_path.exists():
                with open(post_receipt_path) as f:
                    post_data = json.load(f)
                    post_chain_hash = post_data.get("chain_hash", "")

        # Compute chain hash
        chain_hash = self._compute_chain_hash(
            prompt_sha, diff_sha, pre_chain_hash, post_chain_hash, status
        )

        # Truncate diff preview
        diff_preview = all_diffs[:2000] if len(all_diffs) > 2000 else all_diffs

        receipt = {
            # Schema version for forward compatibility
            "schema_version": "1.0.0",

            "receipt_id": receipt_id,
            "receipt_type": "patch_execution",
            "timestamp_utc": timestamp.isoformat(),
            "work_order": "WO-ECHO-PATCH-AGENT-01",
            "prompt": query,
            "prompt_sha256": prompt_sha,
            "action": patch_spec.get("operation", "unknown"),
            "files_modified": patch_spec.get("files", []),
            "file_hashes_before": file_hashes_before,
            "file_hashes_after": file_hashes_after,
            "diff_sha256": diff_sha,
            "diff_preview": diff_preview,
            "lines_added": lines_added,
            "lines_removed": lines_removed,
            "pre_test_receipt": pre_receipt,
            "pre_test_chain_hash": pre_chain_hash,
            "post_test_receipt": post_receipt,
            "post_test_chain_hash": post_chain_hash,
            "backup_manifest": backup_manifest,
            "status": status,
            "chain_hash": chain_hash,

            # Self-documenting verification formulas
            "verification": {
                "chain_hash_formula": "sha256(prompt_sha256 | diff_sha256 | pre_test_chain_hash | post_test_chain_hash | status)",
                "separator": "|",
            },
        }

        if error:
            receipt["error"] = error

        return receipt

    def _save_receipt(self, receipt: Dict[str, Any]) -> str:
        """
        Save receipt to disk.

        Returns relative path to saved receipt.
        """
        self.receipt_dir.mkdir(parents=True, exist_ok=True)

        timestamp = datetime.now(timezone.utc)
        date_str = timestamp.strftime("%Y%m%d")
        receipt_id = receipt.get("receipt_id", generate_receipt_id())[:8]
        status = receipt.get("status", "UNKNOWN").lower()

        filename = f"{status}_{date_str}_{receipt_id}.json"
        receipt_path = self.receipt_dir / filename

        with open(receipt_path, "w") as f:
            json.dump(receipt, f, indent=2, default=str)

        return str(receipt_path.relative_to(PROJECT_ROOT))

    def _rejected_result(
        self, query: str, reason: str, agent_id: str, elapsed_ms: float
    ) -> AgentResult:
        """Build a rejected result."""
        return AgentResult(
            agent_id=agent_id,
            agent_type=self.agent_type,
            query=query,
            result={
                "status": "REJECTED",
                "reason": reason,
            },
            sources=[],
            confidence=0.0,
            elapsed_ms=elapsed_ms,
            error=reason,
        )

    def _rollback_result(
        self,
        query: str,
        post_result: AgentResult,
        backups: Dict[str, str],
        restored_files: List[str],
        agent_id: str,
        elapsed_ms: float,
        receipt_path: str,
    ) -> AgentResult:
        """Build a rollback result."""
        return AgentResult(
            agent_id=agent_id,
            agent_type=self.agent_type,
            query=query,
            result={
                "status": "ROLLBACK",
                "reason": "Post-test failed, changes rolled back",
                "post_test_status": post_result.result.get("status"),
                "post_test_receipt": post_result.result.get("receipt_path"),
                "restored_files": restored_files,
                "receipt_path": receipt_path,
            },
            sources=[f"receipt:{receipt_path}"],
            confidence=0.0,
            elapsed_ms=elapsed_ms,
            error="Post-test failed, changes rolled back",
        )

    async def execute(self, query: str, context: Dict[str, Any]) -> AgentResult:
        """
        Execute a patch with pre/post test gates.

        Flow:
        1. Parse patch intent
        2. Validate file paths against allowlist
        3. Run PRE-TEST via TestAgent
        4. Create backups
        5. Apply patch
        6. Run POST-TEST via TestAgent
        7. If post-test fails, rollback
        8. Generate receipt with chain-hash

        Args:
            query: The user's query
            context: Additional context (patch_spec, etc.)

        Returns:
            AgentResult with patch results and receipt source
        """
        start_time = time.perf_counter()
        agent_id = generate_agent_id()

        if not self.available():
            return AgentResult(
                agent_id=agent_id,
                agent_type=self.agent_type,
                query=query,
                result={"status": "ERROR", "error": "PatchAgent not available"},
                sources=[],
                confidence=0.0,
                elapsed_ms=0.0,
                error="PROJECT_ROOT not writable",
            )

        # 1. Parse patch intent
        patch_spec = self._parse_patch_intent(query, context)
        files = patch_spec.get("files", [])

        if not files:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return self._rejected_result(
                query,
                "No files detected in patch request",
                agent_id,
                elapsed_ms,
            )

        # 2. Validate file paths
        for file_path in files:
            is_valid, reason = self._validate_file_path(file_path)
            if not is_valid:
                elapsed_ms = (time.perf_counter() - start_time) * 1000
                return self._rejected_result(
                    query,
                    f"File path rejected: {file_path} - {reason}",
                    agent_id,
                    elapsed_ms,
                )

        # 3. Run PRE-TEST
        pre_result = await self._run_test_gate("pre")
        pre_status = pre_result.result.get("status")
        pre_receipt_path = pre_result.result.get("receipt_path")

        if pre_status != "PASS":
            elapsed_ms = (time.perf_counter() - start_time) * 1000

            # Build rejection receipt
            receipt = self._build_receipt(
                query=query,
                patch_spec=patch_spec,
                pre_receipt=pre_receipt_path,
                post_receipt=None,
                file_diffs={},
                file_hashes_before={},
                file_hashes_after={},
                backup_manifest={},
                lines_added=0,
                lines_removed=0,
                status="REJECTED",
                error=f"Pre-test failed with status: {pre_status}",
            )
            receipt_path = self._save_receipt(receipt)

            return AgentResult(
                agent_id=agent_id,
                agent_type=self.agent_type,
                query=query,
                result={
                    "status": "REJECTED",
                    "reason": f"Pre-test failed: {pre_status}",
                    "pre_test_receipt": pre_receipt_path,
                    "receipt_path": receipt_path,
                },
                sources=[f"receipt:{receipt_path}"],
                confidence=0.0,
                elapsed_ms=elapsed_ms,
                error=f"Pre-test failed: {pre_status}",
            )

        # 4. Create backups
        file_hashes_before = {f: self._compute_file_hash(f) for f in files}
        backups = self._create_backups(files)

        # 5. Apply patch
        try:
            apply_result = self._apply_patch(patch_spec)
        except Exception as e:
            # Restore backups on error
            self._restore_backups(backups)
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return self._rejected_result(
                query,
                f"Patch failed: {str(e)}",
                agent_id,
                elapsed_ms,
            )

        file_diffs = apply_result["diffs"]
        lines_added = apply_result["lines_added"]
        lines_removed = apply_result["lines_removed"]
        files_modified = apply_result["files_modified"]

        # Get hashes after modification
        file_hashes_after = {f: self._compute_file_hash(f) for f in files}

        # 6. Run POST-TEST
        post_result = await self._run_test_gate("post")
        post_status = post_result.result.get("status")
        post_receipt_path = post_result.result.get("receipt_path")

        if post_status != "PASS":
            # 7. ROLLBACK
            restored_files = self._restore_backups(backups)

            # Build rollback receipt
            receipt = self._build_receipt(
                query=query,
                patch_spec=patch_spec,
                pre_receipt=pre_receipt_path,
                post_receipt=post_receipt_path,
                file_diffs=file_diffs,
                file_hashes_before=file_hashes_before,
                file_hashes_after=file_hashes_after,
                backup_manifest=backups,
                lines_added=lines_added,
                lines_removed=lines_removed,
                status="ROLLBACK",
                error=f"Post-test failed: {post_status}",
            )
            receipt_path = self._save_receipt(receipt)

            elapsed_ms = (time.perf_counter() - start_time) * 1000
            return self._rollback_result(
                query,
                post_result,
                backups,
                restored_files,
                agent_id,
                elapsed_ms,
                receipt_path,
            )

        # 8. Build success receipt
        receipt = self._build_receipt(
            query=query,
            patch_spec=patch_spec,
            pre_receipt=pre_receipt_path,
            post_receipt=post_receipt_path,
            file_diffs=file_diffs,
            file_hashes_before=file_hashes_before,
            file_hashes_after=file_hashes_after,
            backup_manifest=backups,
            lines_added=lines_added,
            lines_removed=lines_removed,
            status="APPLIED",
        )
        receipt_path = self._save_receipt(receipt)

        elapsed_ms = (time.perf_counter() - start_time) * 1000

        return AgentResult(
            agent_id=agent_id,
            agent_type=self.agent_type,
            query=query,
            result={
                "status": "APPLIED",
                "files_modified": files_modified,
                "lines_added": lines_added,
                "lines_removed": lines_removed,
                "pre_receipt": pre_receipt_path,
                "post_receipt": post_receipt_path,
                "receipt_path": receipt_path,
                "chain_hash": receipt.get("chain_hash"),
            },
            sources=[f"receipt:{receipt_path}"],
            confidence=1.0,
            elapsed_ms=elapsed_ms,
            error=None,
        )


# Auto-register the agent when this module is imported
from . import register_agent as _register

try:
    _register(PatchAgent())
except ValueError:
    pass  # Already registered
