"""Gate 5D constrained learned read-only proposals.

The learned controller may choose only deterministic candidate IDs. Candidate
arguments are generated and validated by the runtime before the learned proposal
is considered.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .advisory_arbiter import GuardedAdvisoryArbiter, context_from_instruction
from .bytecode import CognitiveOp, Instruction
from .controller import Controller
from .ledger import canonical_json, stable_hash
from .shadow_controller import ShadowCalibration, ShadowControllerModel, model_input_from_workspace, predict_calibrated_shadow
from .workspace import Workspace
from .workbench import oracle_leak_keys


GATE5D_SCHEMA_VERSION = 1
APPROVED_VERIFIERS = {"internal_verify"}
ALLOWED_LEARNED_ACTIONS = {
    "LIST_FILES": CognitiveOp.CALL_TOOL.value,
    "READ_FILE": CognitiveOp.CALL_TOOL.value,
    "RETRIEVE_MEMORY": CognitiveOp.RETRIEVE_MEMORY.value,
    "INSPECT_WORKSPACE": CognitiveOp.CALL_TOOL.value,
    "RUN_APPROVED_VERIFIER": CognitiveOp.VERIFY.value,
}
FORBIDDEN_LEARNED_OPS = {
    CognitiveOp.PATCH.value,
    CognitiveOp.ANSWER.value,
    CognitiveOp.RETRIEVE_SKILL.value,
    CognitiveOp.CREATE_SKILL.value,
    CognitiveOp.UPDATE_SKILL.value,
    CognitiveOp.STORE_EPISODE.value,
    CognitiveOp.STORE_FAILURE.value,
}


@dataclass(slots=True)
class ReadOnlyCandidateAction:
    candidate_id: str
    action_name: str
    instruction_op: str
    args: dict[str, Any] = field(default_factory=dict)
    risk_class: str = "low"
    source_of_args: str = "runtime_generated"
    observed_path_proof: dict[str, Any] | None = None
    verifier_whitelist_proof: dict[str, Any] | None = None
    legal_under_action_mask: bool = True
    sandbox_contained: bool = True
    mutation_risk: bool = False

    @classmethod
    def create(
        cls,
        action_name: str,
        instruction_op: str,
        args: dict[str, Any] | None = None,
        *,
        risk_class: str = "low",
        source_of_args: str = "runtime_generated",
        observed_path_proof: dict[str, Any] | None = None,
        verifier_whitelist_proof: dict[str, Any] | None = None,
        legal_under_action_mask: bool = True,
        sandbox_contained: bool = True,
        mutation_risk: bool = False,
    ) -> "ReadOnlyCandidateAction":
        payload = {
            "schema_version": GATE5D_SCHEMA_VERSION,
            "action_name": action_name,
            "instruction_op": instruction_op,
            "args": args or {},
            "risk_class": risk_class,
            "source_of_args": source_of_args,
            "observed_path_proof": observed_path_proof,
            "verifier_whitelist_proof": verifier_whitelist_proof,
            "legal_under_action_mask": legal_under_action_mask,
            "sandbox_contained": sandbox_contained,
            "mutation_risk": mutation_risk,
        }
        return cls(candidate_id=_candidate_id(payload), **{key: value for key, value in payload.items() if key != "schema_version"})

    def payload_without_id(self) -> dict[str, Any]:
        return {
            "schema_version": GATE5D_SCHEMA_VERSION,
            "action_name": self.action_name,
            "instruction_op": self.instruction_op,
            "args": self.args,
            "risk_class": self.risk_class,
            "source_of_args": self.source_of_args,
            "observed_path_proof": self.observed_path_proof,
            "verifier_whitelist_proof": self.verifier_whitelist_proof,
            "legal_under_action_mask": self.legal_under_action_mask,
            "sandbox_contained": self.sandbox_contained,
            "mutation_risk": self.mutation_risk,
        }

    def has_stable_id(self) -> bool:
        return self.candidate_id == _candidate_id(self.payload_without_id())

    def to_instruction(self) -> Instruction:
        return Instruction(CognitiveOp(str(self.instruction_op)), args=dict(self.args), rationale=f"Gate 5D approved candidate {self.candidate_id}")

    def to_dict(self) -> dict[str, Any]:
        return {"candidate_id": self.candidate_id, **self.payload_without_id()}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ReadOnlyCandidateAction":
        return cls(
            candidate_id=str(data["candidate_id"]),
            action_name=str(data["action_name"]),
            instruction_op=str(data["instruction_op"]),
            args=dict(data.get("args") or {}),
            risk_class=str(data.get("risk_class", "low")),
            source_of_args=str(data.get("source_of_args", "runtime_generated")),
            observed_path_proof=dict(data["observed_path_proof"]) if isinstance(data.get("observed_path_proof"), dict) else None,
            verifier_whitelist_proof=dict(data["verifier_whitelist_proof"]) if isinstance(data.get("verifier_whitelist_proof"), dict) else None,
            legal_under_action_mask=bool(data.get("legal_under_action_mask", True)),
            sandbox_contained=bool(data.get("sandbox_contained", True)),
            mutation_risk=bool(data.get("mutation_risk", False)),
        )


@dataclass(slots=True)
class LearnedReadOnlyProposal:
    candidate_id: str
    confidence: float
    decision: str
    reason_code: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "candidate_id": self.candidate_id,
            "confidence": self.confidence,
            "decision": self.decision,
            "reason_code": self.reason_code,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "LearnedReadOnlyProposal":
        extra = set(data) - {"candidate_id", "confidence", "decision", "reason_code"}
        if extra:
            raise ValueError("learned proposal may not contain raw args or extra fields: " + ",".join(sorted(extra)))
        return cls(
            candidate_id=str(data.get("candidate_id", "")),
            confidence=float(data.get("confidence", 0.0)),
            decision=str(data.get("decision", "abstain")),
            reason_code=str(data.get("reason_code", "unspecified")),
        )


@dataclass(slots=True)
class ProposalApproval:
    approved: bool
    candidate_id: str
    action_name: str = ""
    reason_code: str = ""
    instruction: Instruction | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "approved": self.approved,
            "candidate_id": self.candidate_id,
            "action_name": self.action_name,
            "reason_code": self.reason_code,
            "instruction": self.instruction.to_dict() if self.instruction is not None else None,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ProposalApproval":
        instruction_data = data.get("instruction")
        return cls(
            approved=bool(data["approved"]),
            candidate_id=str(data.get("candidate_id", "")),
            action_name=str(data.get("action_name", "")),
            reason_code=str(data.get("reason_code", "")),
            instruction=Instruction.from_dict(dict(instruction_data)) if isinstance(instruction_data, dict) else None,
        )


@dataclass(slots=True)
class ProposalAccounting:
    learned_proposals_made: int = 0
    learned_proposals_approved: int = 0
    learned_proposals_rejected: int = 0
    learned_proposals_executed: int = 0
    learned_readonly_actions_executed: int = 0
    learned_verifier_actions_executed: int = 0
    learned_mutating_actions_executed: int = 0
    learned_answer_actions_executed: int = 0
    deterministic_fallback_count: int = 0
    unnecessary_proposal_count: int = 0

    def to_dict(self) -> dict[str, int]:
        return {
            "learned_proposals_made": self.learned_proposals_made,
            "learned_proposals_approved": self.learned_proposals_approved,
            "learned_proposals_rejected": self.learned_proposals_rejected,
            "learned_proposals_executed": self.learned_proposals_executed,
            "learned_readonly_actions_executed": self.learned_readonly_actions_executed,
            "learned_verifier_actions_executed": self.learned_verifier_actions_executed,
            "learned_mutating_actions_executed": self.learned_mutating_actions_executed,
            "learned_answer_actions_executed": self.learned_answer_actions_executed,
            "deterministic_fallback_count": self.deterministic_fallback_count,
            "unnecessary_proposal_count": self.unnecessary_proposal_count,
        }

    def merge(self, other: "ProposalAccounting") -> None:
        for key, value in other.to_dict().items():
            setattr(self, key, int(getattr(self, key)) + int(value))


@dataclass(slots=True)
class ReadOnlyProof:
    readonly_action_count: int = 0
    verifier_action_count: int = 0
    mutation_detected_count: int = 0
    file_write_detected_count: int = 0
    memory_write_detected_count: int = 0
    skill_write_detected_count: int = 0
    answer_write_detected_count: int = 0
    patch_generated_count: int = 0
    shell_command_generated_count: int = 0
    audit_log_append_count: int = 0

    def record_candidate(self, candidate: ReadOnlyCandidateAction) -> None:
        text = canonical_json(candidate.to_dict()).lower()
        if candidate.mutation_risk or candidate.instruction_op in FORBIDDEN_LEARNED_OPS:
            self.mutation_detected_count += 1
        elif candidate.action_name == "RUN_APPROVED_VERIFIER":
            self.verifier_action_count += 1
        else:
            self.readonly_action_count += 1
        if candidate.instruction_op == CognitiveOp.PATCH.value or "patch" in text:
            self.patch_generated_count += 1
            self.file_write_detected_count += 1
        if candidate.instruction_op == CognitiveOp.ANSWER.value:
            self.answer_write_detected_count += 1
        if _contains_shell_text(text):
            self.shell_command_generated_count += 1

    def to_dict(self) -> dict[str, int]:
        return {
            "readonly_action_count": self.readonly_action_count,
            "verifier_action_count": self.verifier_action_count,
            "mutation_detected_count": self.mutation_detected_count,
            "file_write_detected_count": self.file_write_detected_count,
            "memory_write_detected_count": self.memory_write_detected_count,
            "skill_write_detected_count": self.skill_write_detected_count,
            "answer_write_detected_count": self.answer_write_detected_count,
            "patch_generated_count": self.patch_generated_count,
            "shell_command_generated_count": self.shell_command_generated_count,
            "audit_log_append_count": self.audit_log_append_count,
        }

    def merge(self, other: "ReadOnlyProof") -> None:
        for key, value in other.to_dict().items():
            setattr(self, key, int(getattr(self, key)) + int(value))


class CandidateActionGenerator:
    def generate(
        self,
        workspace: Workspace,
        model_input: dict[str, Any],
        action_mask: dict[str, bool],
        sandbox_root: str | Path | None = None,
    ) -> list[ReadOnlyCandidateAction]:
        _reject_oracle_leaks(model_input)
        candidates: list[ReadOnlyCandidateAction] = []
        stage = str(workspace.metadata.get("stage", "created"))
        if stage == "goal_created":
            candidates.append(
                self._candidate(
                    "RETRIEVE_MEMORY",
                    CognitiveOp.RETRIEVE_MEMORY.value,
                    {"query": workspace.task},
                    "workspace_task",
                    action_mask,
                )
            )
        if stage == "hypothesis_created":
            candidates.append(
                self._candidate(
                    "RUN_APPROVED_VERIFIER",
                    CognitiveOp.VERIFY.value,
                    {"verifier": "internal_verify"},
                    "verifier_whitelist",
                    action_mask,
                    verifier_whitelist_proof={"verifier": "internal_verify", "allowed": True},
                )
            )
        if "inspect_workspace" in canonical_json(model_input).lower():
            candidates.append(
                self._candidate(
                    "INSPECT_WORKSPACE",
                    CognitiveOp.CALL_TOOL.value,
                    {"operation": "inspect_workspace"},
                    "workspace_summary",
                    action_mask,
                )
            )
        observed_paths = _observed_paths(workspace, model_input)
        if sandbox_root is not None and observed_paths:
            root = Path(sandbox_root)
            list_contained = _sandbox_contained(root, ".")
            candidates.append(
                self._candidate(
                    "LIST_FILES",
                    CognitiveOp.CALL_TOOL.value,
                    {"operation": "list_files", "root": "sandbox"},
                    "observed_sandbox",
                    action_mask,
                    sandbox_contained=list_contained,
                )
            )
            for path in sorted(observed_paths):
                if not _safe_relative_path(path):
                    continue
                contained = _sandbox_contained(root, path)
                if not contained:
                    continue
                candidates.append(
                    self._candidate(
                        "READ_FILE",
                        CognitiveOp.CALL_TOOL.value,
                        {"operation": "read_file", "path": path},
                        "observed_file_listing",
                        action_mask,
                        observed_path_proof={"path": path, "observed": True},
                        sandbox_contained=True,
                    )
                )
        return [candidate for candidate in candidates if validate_candidate(candidate)[0]]

    def _candidate(
        self,
        action_name: str,
        instruction_op: str,
        args: dict[str, Any],
        source_of_args: str,
        action_mask: dict[str, bool],
        *,
        observed_path_proof: dict[str, Any] | None = None,
        verifier_whitelist_proof: dict[str, Any] | None = None,
        sandbox_contained: bool = True,
    ) -> ReadOnlyCandidateAction:
        return ReadOnlyCandidateAction.create(
            action_name=action_name,
            instruction_op=instruction_op,
            args=args,
            risk_class="low",
            source_of_args=source_of_args,
            observed_path_proof=observed_path_proof,
            verifier_whitelist_proof=verifier_whitelist_proof,
            legal_under_action_mask=bool(action_mask.get(instruction_op, False) or action_mask.get(action_name, False)),
            sandbox_contained=sandbox_contained,
            mutation_risk=False,
        )


class LearnedProposalSelector:
    def propose(
        self,
        model: ShadowControllerModel,
        calibration: ShadowCalibration,
        model_input: dict[str, Any],
        candidates: list[ReadOnlyCandidateAction],
    ) -> LearnedReadOnlyProposal:
        if not candidates:
            return LearnedReadOnlyProposal("", 0.0, "abstain", "no_candidates")
        prediction = predict_calibrated_shadow(model, calibration, model_input)
        if prediction.abstained:
            return LearnedReadOnlyProposal("", prediction.confidence, "abstain", prediction.abstention_reason)
        ordered = sorted(
            candidates,
            key=lambda candidate: (
                candidate.instruction_op != prediction.action,
                candidate.action_name != "RETRIEVE_MEMORY",
                candidate.candidate_id,
            ),
        )
        selected = ordered[0]
        if prediction.confidence < calibration.config.confidence_threshold:
            return LearnedReadOnlyProposal("", prediction.confidence, "abstain", "confidence_below_threshold")
        return LearnedReadOnlyProposal(selected.candidate_id, prediction.confidence, "accept", "selected_from_deterministic_candidates")


class ReadOnlyProposalApprover:
    def approve(
        self,
        proposal: LearnedReadOnlyProposal,
        candidates: list[ReadOnlyCandidateAction],
        action_mask: dict[str, bool],
        advisory_arbiter: GuardedAdvisoryArbiter,
    ) -> ProposalApproval:
        if proposal.decision != "accept":
            return ProposalApproval(False, proposal.candidate_id, reason_code=proposal.reason_code)
        by_id = {candidate.candidate_id: candidate for candidate in candidates}
        candidate = by_id.get(proposal.candidate_id)
        if candidate is None:
            return ProposalApproval(False, proposal.candidate_id, reason_code="candidate_id_not_found")
        valid, reason = validate_candidate(candidate)
        if not valid:
            return ProposalApproval(False, candidate.candidate_id, candidate.action_name, reason)
        if not bool(action_mask.get(candidate.instruction_op, False) or action_mask.get(candidate.action_name, False)):
            return ProposalApproval(False, candidate.candidate_id, candidate.action_name, "action_mask_violation")
        context = context_from_instruction(
            candidate.to_instruction(),
            verifier_context="verifier_status=not_run",
            action_risk_class="low",
        )
        advisory = advisory_arbiter.decide(context)
        if not advisory.accepted_by_arbiter:
            return ProposalApproval(False, candidate.candidate_id, candidate.action_name, "advisory_blocked:" + advisory.reason)
        return ProposalApproval(True, candidate.candidate_id, candidate.action_name, "approved", candidate.to_instruction())


class LearnedReadOnlyProposalController:
    def __init__(
        self,
        base: Controller,
        model: ShadowControllerModel,
        calibration: ShadowCalibration,
        *,
        generator: CandidateActionGenerator | None = None,
        selector: LearnedProposalSelector | None = None,
        approver: ReadOnlyProposalApprover | None = None,
        advisory_arbiter: GuardedAdvisoryArbiter | None = None,
        sandbox_root: str | Path | None = None,
    ) -> None:
        self.base = base
        self.model = model
        self.calibration = calibration
        self.generator = generator or CandidateActionGenerator()
        self.selector = selector or LearnedProposalSelector()
        self.approver = approver or ReadOnlyProposalApprover()
        self.advisory_arbiter = advisory_arbiter or GuardedAdvisoryArbiter()
        self.sandbox_root = Path(sandbox_root) if sandbox_root is not None else None
        self.accounting = ProposalAccounting()
        self.readonly_proof = ReadOnlyProof()
        self.logs: list[dict[str, Any]] = []

    def next_instruction(self, workspace: Workspace) -> Instruction:
        deterministic = self.base.next_instruction(workspace)
        model_input = model_input_from_workspace(workspace, previous_action=self.logs[-1]["executed_action"] if self.logs else None)
        action_mask = dict(model_input.get("action_mask") or {})
        candidates = self.generator.generate(workspace, model_input, action_mask, self.sandbox_root)
        equivalent = [candidate for candidate in candidates if candidate.instruction_op == deterministic.op.value]
        proposal = self.selector.propose(self.model, self.calibration, model_input, equivalent or candidates)
        approval = self.approver.approve(proposal, candidates, action_mask, self.advisory_arbiter)
        executed = deterministic
        executed_by_learned = False
        if proposal.candidate_id:
            self.accounting.learned_proposals_made += 1
        if approval.approved and approval.instruction is not None:
            self.accounting.learned_proposals_approved += 1
            self.accounting.learned_proposals_executed += 1
            executed = approval.instruction
            executed_by_learned = True
            candidate = next(item for item in candidates if item.candidate_id == approval.candidate_id)
            self.readonly_proof.record_candidate(candidate)
            if candidate.action_name == "RUN_APPROVED_VERIFIER":
                self.accounting.learned_verifier_actions_executed += 1
            else:
                self.accounting.learned_readonly_actions_executed += 1
            if executed.op != deterministic.op:
                self.accounting.unnecessary_proposal_count += 1
        else:
            if proposal.candidate_id:
                self.accounting.learned_proposals_rejected += 1
                self.accounting.deterministic_fallback_count += 1
        if executed.op in {CognitiveOp.PATCH, CognitiveOp.STORE_EPISODE, CognitiveOp.STORE_FAILURE, CognitiveOp.CREATE_SKILL, CognitiveOp.UPDATE_SKILL}:
            self.accounting.learned_mutating_actions_executed += int(executed_by_learned)
        if executed.op == CognitiveOp.ANSWER:
            self.accounting.learned_answer_actions_executed += int(executed_by_learned)
        self.logs.append(
            {
                "step_index": len(self.logs),
                "deterministic_action": deterministic.op.value,
                "executed_action": executed.op.value,
                "executed_by_learned_proposal": executed_by_learned,
                "proposal": proposal.to_dict(),
                "approval": approval.to_dict(),
                "candidate_count": len(candidates),
                "candidate_ids": [candidate.candidate_id for candidate in candidates],
            }
        )
        return executed

    def summary(self) -> dict[str, Any]:
        return {
            "enabled": True,
            "no_authoritative_mutating_learned_control": True,
            "proposal_accounting": self.accounting.to_dict(),
            "readonly_proof": self.readonly_proof.to_dict(),
            "logs": self.logs,
        }


def validate_candidate(candidate: ReadOnlyCandidateAction) -> tuple[bool, str]:
    if not candidate.has_stable_id():
        return False, "candidate_id_not_deterministic"
    if candidate.action_name not in ALLOWED_LEARNED_ACTIONS:
        return False, "illegal_learned_action"
    if candidate.instruction_op != ALLOWED_LEARNED_ACTIONS[candidate.action_name]:
        return False, "instruction_mapping_mismatch"
    if candidate.instruction_op in FORBIDDEN_LEARNED_OPS:
        return False, "forbidden_instruction_op"
    if candidate.mutation_risk:
        return False, "mutation_risk"
    if not candidate.legal_under_action_mask:
        return False, "action_mask_violation"
    if not candidate.sandbox_contained:
        return False, "sandbox_not_contained"
    text = canonical_json(candidate.to_dict())
    if _contains_oracle_or_canary(text) or oracle_leak_keys(candidate.to_dict()):
        return False, "oracle_or_canary_leak"
    if _contains_shell_text(text):
        return False, "shell_generation"
    if candidate.action_name == "READ_FILE":
        path = str(candidate.args.get("path", ""))
        if not _safe_relative_path(path):
            return False, "unsafe_path"
        if not candidate.observed_path_proof or candidate.observed_path_proof.get("path") != path:
            return False, "unobserved_path"
    if candidate.action_name == "RUN_APPROVED_VERIFIER":
        verifier = str(candidate.args.get("verifier", ""))
        if verifier not in APPROVED_VERIFIERS:
            return False, "illegal_verifier"
        if not candidate.verifier_whitelist_proof or not candidate.verifier_whitelist_proof.get("allowed"):
            return False, "missing_verifier_whitelist_proof"
    return True, "ok"


def proposal_artifact_hash(payload: dict[str, Any]) -> str:
    return stable_hash({"schema_version": GATE5D_SCHEMA_VERSION, "payload": payload})


def _candidate_id(payload: dict[str, Any]) -> str:
    return "c_" + stable_hash(payload)[:16]


def _reject_oracle_leaks(model_input: dict[str, Any]) -> None:
    text = canonical_json(model_input)
    if _contains_oracle_or_canary(text) or oracle_leak_keys(model_input):
        raise ValueError("candidate generation received hidden oracle or canary data")


def _contains_oracle_or_canary(text: str) -> bool:
    lowered = text.lower()
    return "canary_" in lowered or "expected_fix" in lowered or "hidden_verifier_oracle" in lowered or "defect_kind" in lowered or "task_family" in lowered


def _contains_shell_text(text: str) -> bool:
    lowered = text.lower()
    markers = ["python -c", "bash -lc", "cmd /c", "powershell", "rm -rf", "shell", "$("]
    return any(marker in lowered for marker in markers)


def _safe_relative_path(path: str) -> bool:
    if not path:
        return False
    requested = Path(path)
    return not requested.is_absolute() and ".." not in requested.parts and not _contains_shell_text(path)


def _sandbox_contained(root: Path, path: str) -> bool:
    try:
        requested = Path(path)
        if requested == Path("."):
            target = root.resolve()
        else:
            if not _safe_relative_path(path):
                return False
            target = (root / requested).resolve()
        return target == root.resolve() or target.is_relative_to(root.resolve())
    except OSError:
        return False


def _observed_paths(workspace: Workspace, model_input: dict[str, Any]) -> set[str]:
    paths: set[str] = set()
    for item in workspace.metadata.get("observed_paths", []) or []:
        paths.add(str(item))
    sandbox_files = workspace.metadata.get("sandbox_files")
    if isinstance(sandbox_files, dict):
        paths.update(str(path) for path in sandbox_files)
    observation = dict(model_input.get("agent_visible_observation") or {})
    for key in ["observed_paths", "visible_file_paths"]:
        for item in observation.get(key, []) or []:
            paths.add(str(item))
    for summary in observation.get("visible_tool_summaries", []) or []:
        if not isinstance(summary, dict):
            continue
        if "path" in summary:
            paths.add(str(summary["path"]))
        for item in summary.get("files", []) or []:
            paths.add(str(item))
    return {path for path in paths if _safe_relative_path(path)}
