"""Deterministic application of allowlisted static repair operations."""

from __future__ import annotations

import copy
import hashlib
import json
from typing import Any

from .features import STATEFUL_SHELL_RE
from .io import get_sample_id
from .policy import PolicyViolation, enforce_repair_policy


def _hash(value: Any) -> str:
    """Return a stable SHA-256 digest for audit records."""

    encoded = json.dumps(value, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def _assistant_command(message: dict[str, Any]) -> str:
    """Extract a shell command from the first tool call in an assistant message."""

    calls = message.get("tool_calls") or []
    if not calls or not isinstance(calls[0], dict):
        return ""
    function = calls[0].get("function") or {}
    raw_arguments = function.get("arguments", "{}") if isinstance(function, dict) else "{}"
    try:
        arguments = json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
    except json.JSONDecodeError:
        return ""
    return str(arguments.get("command") or arguments.get("cmd") or "") if isinstance(arguments, dict) else ""


def _is_read_only_pair(assistant_message: dict[str, Any], tool_message: dict[str, Any]) -> bool:
    """Conservatively determine whether a tool pair has no repository side effects."""

    if assistant_message.get("role") != "assistant" or tool_message.get("role") != "tool":
        return False
    calls = assistant_message.get("tool_calls") or []
    if len(calls) != 1 or not isinstance(calls[0], dict):
        return False
    function = calls[0].get("function") or {}
    name = function.get("name") if isinstance(function, dict) else None
    if name == "str_replace_editor":
        raw_arguments = function.get("arguments", "{}")
        try:
            arguments = json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
        except json.JSONDecodeError:
            return False
        return isinstance(arguments, dict) and arguments.get("command") == "view"
    command = _assistant_command(assistant_message)
    return bool(command) and not STATEFUL_SHELL_RE.search(command)


def _replace_message(
    entry: dict[str, Any],
    operation: str,
    replacement: Any,
) -> None:
    """Apply a non-deleting operation to one annotated trajectory entry."""

    message = entry["message"]
    if message.get("role") == "tool":
        raise PolicyViolation(f"{operation} cannot modify a tool observation")

    if operation in {"REMOVE_DUPLICATE_PROSE", "REWRITE_FINAL_SUMMARY", "REDACT_SECRET"}:
        if not isinstance(replacement, str):
            raise PolicyViolation(f"{operation} requires a string replacement")
        if message.get("role") != "assistant":
            raise PolicyViolation(f"{operation} is restricted to assistant messages")
        message["content"] = replacement
        return

    if operation in {
        "CANONICALIZE_JSON",
        "NORMALIZE_TOOL_ALIAS",
        "EMIT_STRUCTURALLY_CORRECT_ASSISTANT_TURN",
    }:
        if not isinstance(replacement, dict) or replacement.get("role") != "assistant":
            raise PolicyViolation(f"{operation} requires a complete assistant-message object")
        entry["message"] = copy.deepcopy(replacement)
        return

    raise PolicyViolation(f"Unsupported non-deleting operation: {operation}")


def apply_static_repair(
    record: dict[str, Any],
    repair_plan: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Apply an approved allowlisted plan and return a structured audit diff.

    The function performs only syntactic operations. It never executes code and
    never changes outcome labels or patch metadata.
    """

    enforce_repair_policy(repair_plan)
    decision = repair_plan["repair_decision"]
    if decision not in {"APPLY_STATIC_REPAIR", "CREATE_STEP_EXAMPLE"}:
        raise PolicyViolation(f"Repair decision {decision} does not authorize mutation")

    original = copy.deepcopy(record)
    repaired = copy.deepcopy(record)
    entries = [
        {"original_turn_id": turn_id, "message": copy.deepcopy(message)}
        for turn_id, message in enumerate(repaired.get("trajectory") or [], 1)
    ]
    removed_turns: list[int] = []
    replaced_turns: list[int] = []

    def find_entry(turn_id: int) -> dict[str, Any]:
        for item in entries:
            if item["original_turn_id"] == turn_id:
                return item
        raise PolicyViolation(f"Repair plan references missing or already removed turn {turn_id}")

    for operation in repair_plan["operations"]:
        op = operation["op"]
        targets = operation["target_turns"]
        replacement = operation["replacement"]

        if op == "TRUNCATE_BEFORE_TURN":
            if len(targets) != 1:
                raise PolicyViolation("TRUNCATE_BEFORE_TURN requires exactly one target turn")
            cutoff = targets[0]
            to_remove = [item["original_turn_id"] for item in entries if item["original_turn_id"] >= cutoff]
            entries[:] = [item for item in entries if item["original_turn_id"] < cutoff]
            removed_turns.extend(to_remove)
            continue

        if op == "DROP_REDUNDANT_READ_ONLY_PAIR":
            if len(targets) != 2 or targets[1] != targets[0] + 1:
                raise PolicyViolation("A read-only pair must contain two consecutive original turns")
            assistant_entry = find_entry(targets[0])
            tool_entry = find_entry(targets[1])
            if not _is_read_only_pair(assistant_entry["message"], tool_entry["message"]):
                raise PolicyViolation("The requested pair is not provably read-only")
            entries[:] = [item for item in entries if item["original_turn_id"] not in set(targets)]
            removed_turns.extend(targets)
            continue

        if len(targets) != 1:
            raise PolicyViolation(f"{op} requires exactly one target turn")
        entry = find_entry(targets[0])
        _replace_message(entry, op, replacement)
        replaced_turns.append(targets[0])

    repaired["trajectory"] = [item["message"] for item in entries]

    # Deterministic invariants protect immutable outcome and patch metadata.
    if repaired.get("resolved") != original.get("resolved"):
        raise PolicyViolation("Static repair changed resolved")
    original_metadata = original.get("metadata") or {}
    repaired_metadata = repaired.get("metadata") or {}
    for key in ("model_patch", "reference_patch"):
        if repaired_metadata.get(key) != original_metadata.get(key):
            raise PolicyViolation(f"Static repair changed immutable metadata field {key}")

    # Every retained tool observation must remain byte-for-byte equivalent.
    original_tool_hashes = {
        turn_id: _hash(message)
        for turn_id, message in enumerate(original.get("trajectory") or [], 1)
        if isinstance(message, dict) and message.get("role") == "tool"
    }
    retained_entries = {item["original_turn_id"]: item["message"] for item in entries}
    for turn_id, digest in original_tool_hashes.items():
        if turn_id in retained_entries and _hash(retained_entries[turn_id]) != digest:
            raise PolicyViolation(f"Static repair modified retained tool output at turn {turn_id}")

    diff = {
        "sample_id": get_sample_id(record),
        "input_sha256": _hash(original),
        "output_sha256": _hash(repaired),
        "removed_original_turn_ids": sorted(set(removed_turns)),
        "replaced_original_turn_ids": sorted(set(replaced_turns)),
        "original_turn_count": len(original.get("trajectory") or []),
        "repaired_turn_count": len(repaired.get("trajectory") or []),
        "maximum_training_use": repair_plan["maximum_training_use"],
    }
    return repaired, diff