"""Gate 7A dry-run learned memory write recommendations.

The learned path may recommend only deterministic candidate IDs. Candidate
memory records are built by deterministic runtime code from grounded trace
evidence, and approvals are logged as dry-run records without writing to the
authoritative memory store.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .bytecode import CognitiveOp
from .consolidation import ConsolidatedMemoryRecord, SUPPORTED_MEMORY_TYPES
from .ledger import TraceLedgerRecord, canonical_json, stable_hash
from .workbench import oracle_leak_keys


GATE7A_SCHEMA_VERSION = 1
MEMORY_WRITE_CANDIDATE_TYPE = "MEMORY_WRITE_CANDIDATE"
MEMORY_WRITE_RECOMMENDATION_TYPE = "MEMORY_WRITE_RECOMMENDATION"
DETERMINISTIC_MEMORY_CANDIDATE_GENERATOR = "deterministic_memory_candidate_generator_v1"
FORBIDDEN_LEARNED_MEMORY_FIELDS = {
    "claim",
    "memory",
    "memory_record",
    "memory_type",
    "metadata",
    "source_trace_hashes",
    "source_trace_ids",
    "verifier_evidence",
    "write",
    "write_text",
}


@dataclass(slots=True)
class MemoryWriteCandidate:
    candidate_id: str
    candidate_type: str
    memory_type: str
    claim: str
    source_trace_ids: list[str]
    source_trace_hashes: list[str]
    verifier_evidence: list[dict[str, Any]]
    confidence: float
    status: str = "candidate"
    agent_safe: bool = True
    metadata: dict[str, Any] = field(default_factory=dict)
    dry_run_only: bool = True
    evidence_required: bool = True
    deterministic_generator: str = DETERMINISTIC_MEMORY_CANDIDATE_GENERATOR
    legal_under_action_mask: bool = True
    oracle_free: bool = True

    @classmethod
    def create(
        cls,
        *,
        memory_type: str,
        claim: str,
        traces: list[TraceLedgerRecord],
        verifier_evidence: list[dict[str, Any]] | None = None,
        confidence: float = 1.0,
        metadata: dict[str, Any] | None = None,
        agent_safe: bool = True,
        dry_run_only: bool = True,
        evidence_required: bool = True,
        legal_under_action_mask: bool = True,
        oracle_free: bool = True,
    ) -> "MemoryWriteCandidate":
        trace_ids = [trace.trace_id for trace in traces]
        trace_hashes = [trace.content_hash for trace in traces]
        evidence = [dict(item) for item in (verifier_evidence or _verifier_evidence_from_traces(traces))]
        payload = {
            "schema_version": GATE7A_SCHEMA_VERSION,
            "candidate_type": MEMORY_WRITE_CANDIDATE_TYPE,
            "memory_type": memory_type,
            "claim": claim,
            "source_trace_ids": trace_ids,
            "source_trace_hashes": trace_hashes,
            "verifier_evidence": evidence,
            "confidence": float(confidence),
            "status": "candidate",
            "agent_safe": bool(agent_safe),
            "metadata": dict(metadata or {}),
            "dry_run_only": bool(dry_run_only),
            "evidence_required": bool(evidence_required),
            "deterministic_generator": DETERMINISTIC_MEMORY_CANDIDATE_GENERATOR,
            "legal_under_action_mask": bool(legal_under_action_mask),
            "oracle_free": bool(oracle_free),
        }
        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": GATE7A_SCHEMA_VERSION,
            "candidate_type": self.candidate_type,
            "memory_type": self.memory_type,
            "claim": self.claim,
            "source_trace_ids": list(self.source_trace_ids),
            "source_trace_hashes": list(self.source_trace_hashes),
            "verifier_evidence": [dict(item) for item in self.verifier_evidence],
            "confidence": self.confidence,
            "status": self.status,
            "agent_safe": self.agent_safe,
            "metadata": dict(self.metadata),
            "dry_run_only": self.dry_run_only,
            "evidence_required": self.evidence_required,
            "deterministic_generator": self.deterministic_generator,
            "legal_under_action_mask": self.legal_under_action_mask,
            "oracle_free": self.oracle_free,
        }

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

    def to_memory_record(self) -> ConsolidatedMemoryRecord:
        return ConsolidatedMemoryRecord(
            memory_id="mem_" + stable_hash(self.payload_without_id())[:24],
            memory_type=self.memory_type,
            claim=self.claim,
            source_trace_ids=list(self.source_trace_ids),
            source_trace_hashes=list(self.source_trace_hashes),
            verifier_evidence=[dict(item) for item in self.verifier_evidence],
            confidence=self.confidence,
            status="active",
            created_by="gate4_consolidator",
            agent_safe=self.agent_safe,
            metadata=dict(self.metadata),
        )

    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]) -> "MemoryWriteCandidate":
        return cls(
            candidate_id=str(data["candidate_id"]),
            candidate_type=str(data.get("candidate_type", MEMORY_WRITE_CANDIDATE_TYPE)),
            memory_type=str(data["memory_type"]),
            claim=str(data.get("claim", "")),
            source_trace_ids=[str(item) for item in data.get("source_trace_ids", [])],
            source_trace_hashes=[str(item) for item in data.get("source_trace_hashes", [])],
            verifier_evidence=[dict(item) for item in data.get("verifier_evidence", [])],
            confidence=float(data.get("confidence", 0.0)),
            status=str(data.get("status", "candidate")),
            agent_safe=bool(data.get("agent_safe", True)),
            metadata=dict(data.get("metadata") or {}),
            dry_run_only=bool(data.get("dry_run_only", True)),
            evidence_required=bool(data.get("evidence_required", True)),
            deterministic_generator=str(data.get("deterministic_generator", DETERMINISTIC_MEMORY_CANDIDATE_GENERATOR)),
            legal_under_action_mask=bool(data.get("legal_under_action_mask", True)),
            oracle_free=bool(data.get("oracle_free", True)),
        )


@dataclass(slots=True)
class LearnedMemoryRecommendation:
    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]) -> "LearnedMemoryRecommendation":
        extra = set(data) - {"candidate_id", "confidence", "decision", "reason_code"}
        if extra:
            forbidden = sorted(extra | (set(data) & FORBIDDEN_LEARNED_MEMORY_FIELDS))
            raise ValueError("learned memory recommendation may not contain memory content or extra fields: " + ",".join(forbidden))
        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 MemoryRecommendationApproval:
    approved: bool
    candidate_id: str
    reason_code: str = ""
    dry_run_record: dict[str, Any] | None = None

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

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


@dataclass(slots=True)
class MemoryRecommendationAccounting:
    learned_memory_recommendations_made: int = 0
    learned_memory_recommendations_approved: int = 0
    learned_memory_recommendations_rejected: int = 0
    learned_memory_recommendations_logged: int = 0
    learned_memory_writes_executed: int = 0
    authoritative_memory_write_count: int = 0
    memory_store_write_count: int = 0
    file_write_count: int = 0
    deterministic_fallback_count: int = 0

    def to_dict(self) -> dict[str, int]:
        return {
            "learned_memory_recommendations_made": self.learned_memory_recommendations_made,
            "learned_memory_recommendations_approved": self.learned_memory_recommendations_approved,
            "learned_memory_recommendations_rejected": self.learned_memory_recommendations_rejected,
            "learned_memory_recommendations_logged": self.learned_memory_recommendations_logged,
            "learned_memory_writes_executed": self.learned_memory_writes_executed,
            "authoritative_memory_write_count": self.authoritative_memory_write_count,
            "memory_store_write_count": self.memory_store_write_count,
            "file_write_count": self.file_write_count,
            "deterministic_fallback_count": self.deterministic_fallback_count,
        }


@dataclass(slots=True)
class MemoryRecommendationDryRunProof:
    dry_run_recommendation_count: int = 0
    learned_memory_writes_executed: int = 0
    authoritative_memory_write_count: int = 0
    memory_store_write_count: int = 0
    file_write_count: int = 0
    learned_memory_content_generated_count: int = 0
    audit_log_append_count: int = 0

    def record_dry_run(self, recommendation: LearnedMemoryRecommendation, candidate: MemoryWriteCandidate) -> None:
        self.dry_run_recommendation_count += 1
        self.audit_log_append_count += 1
        recommendation_payload = recommendation.to_dict()
        if any(key in recommendation_payload for key in FORBIDDEN_LEARNED_MEMORY_FIELDS):
            self.learned_memory_content_generated_count += 1
        if not candidate.dry_run_only:
            self.learned_memory_writes_executed += 1

    def record_forbidden_attempt(self, kind: str) -> None:
        if kind == "memory_write":
            self.learned_memory_writes_executed += 1
        elif kind == "authoritative_memory_write":
            self.authoritative_memory_write_count += 1
        elif kind == "memory_store_write":
            self.memory_store_write_count += 1
        elif kind == "file_write":
            self.file_write_count += 1
        elif kind == "memory_content":
            self.learned_memory_content_generated_count += 1

    def to_dict(self) -> dict[str, int]:
        return {
            "dry_run_recommendation_count": self.dry_run_recommendation_count,
            "learned_memory_writes_executed": self.learned_memory_writes_executed,
            "authoritative_memory_write_count": self.authoritative_memory_write_count,
            "memory_store_write_count": self.memory_store_write_count,
            "file_write_count": self.file_write_count,
            "learned_memory_content_generated_count": self.learned_memory_content_generated_count,
            "audit_log_append_count": self.audit_log_append_count,
        }


class MemoryWriteCandidateGenerator:
    def generate(
        self,
        traces: list[TraceLedgerRecord],
        *,
        opportunity_id: str,
        memory_type: str = "skill_evidence",
        claim: str | None = None,
        metadata: dict[str, Any] | None = None,
        action_mask: dict[str, bool] | None = None,
        include_unsupported: bool = False,
    ) -> list[MemoryWriteCandidate]:
        if not traces:
            return []
        mask = action_mask or {CognitiveOp.STORE_EPISODE.value: True}
        base_claim = claim or f"Verified trace bundle `{opportunity_id}` is eligible for dry-run memory recommendation."
        candidate = MemoryWriteCandidate.create(
            memory_type=memory_type,
            claim=base_claim,
            traces=traces,
            metadata={"opportunity_id": opportunity_id, **dict(metadata or {})},
            legal_under_action_mask=bool(mask.get(CognitiveOp.STORE_EPISODE.value, False)),
        )
        candidates = [candidate]
        if include_unsupported:
            candidates.append(
                MemoryWriteCandidate.create(
                    memory_type="unsupported_autonomous_memory",
                    claim=f"Unsupported memory probe `{opportunity_id}` must be rejected.",
                    traces=traces,
                    metadata={"opportunity_id": f"{opportunity_id}_unsupported"},
                    legal_under_action_mask=bool(mask.get(CognitiveOp.STORE_EPISODE.value, False)),
                )
            )
        return _dedupe_candidates(candidates)


class LearnedMemoryRecommendationSelector:
    def propose(self, candidates: list[MemoryWriteCandidate]) -> LearnedMemoryRecommendation:
        safe_candidates = [candidate for candidate in candidates if validate_memory_write_candidate(candidate)[0]]
        if not safe_candidates:
            return LearnedMemoryRecommendation("", 0.0, "abstain", "no_valid_memory_candidates")
        ordered = sorted(
            safe_candidates,
            key=lambda candidate: (-len(candidate.verifier_evidence), -candidate.confidence, candidate.memory_type, candidate.candidate_id),
        )
        return LearnedMemoryRecommendation(ordered[0].candidate_id, 0.94, "accept", "selected_from_deterministic_memory_candidates")


class MemoryRecommendationArbiter:
    def approve(
        self,
        recommendation: LearnedMemoryRecommendation,
        candidates: list[MemoryWriteCandidate],
        action_mask: dict[str, bool] | None = None,
    ) -> MemoryRecommendationApproval:
        if recommendation.decision != "accept":
            return MemoryRecommendationApproval(False, recommendation.candidate_id, recommendation.reason_code)
        by_id = {candidate.candidate_id: candidate for candidate in candidates}
        candidate = by_id.get(recommendation.candidate_id)
        if candidate is None:
            return MemoryRecommendationApproval(False, recommendation.candidate_id, "candidate_id_not_found")
        valid, reason = validate_memory_write_candidate(candidate)
        if not valid:
            return MemoryRecommendationApproval(False, candidate.candidate_id, reason)
        mask = action_mask or {CognitiveOp.STORE_EPISODE.value: True}
        if not bool(mask.get(CognitiveOp.STORE_EPISODE.value, False)):
            return MemoryRecommendationApproval(False, candidate.candidate_id, "action_mask_violation")
        return MemoryRecommendationApproval(
            True,
            candidate.candidate_id,
            "approved_dry_run_only",
            {
                "candidate_id": candidate.candidate_id,
                "memory_type": candidate.memory_type,
                "dry_run_only": True,
                "evidence_required": candidate.evidence_required,
                "source_trace_count": len(candidate.source_trace_ids),
            },
        )


class LearnedMemoryRecommendationController:
    def __init__(
        self,
        *,
        selector: LearnedMemoryRecommendationSelector | None = None,
        arbiter: MemoryRecommendationArbiter | None = None,
    ) -> None:
        self.selector = selector or LearnedMemoryRecommendationSelector()
        self.arbiter = arbiter or MemoryRecommendationArbiter()
        self.accounting = MemoryRecommendationAccounting()
        self.dry_run_proof = MemoryRecommendationDryRunProof()
        self.logs: list[dict[str, Any]] = []

    def review(
        self,
        candidates: list[MemoryWriteCandidate],
        *,
        opportunity_id: str,
        action_mask: dict[str, bool] | None = None,
    ) -> dict[str, Any]:
        recommendation = self.selector.propose(candidates)
        approval = self.arbiter.approve(recommendation, candidates, action_mask)
        if recommendation.candidate_id:
            self.accounting.learned_memory_recommendations_made += 1
        if approval.approved:
            self.accounting.learned_memory_recommendations_approved += 1
            self.accounting.learned_memory_recommendations_logged += 1
            candidate = next(item for item in candidates if item.candidate_id == approval.candidate_id)
            self.dry_run_proof.record_dry_run(recommendation, candidate)
        elif recommendation.candidate_id:
            self.accounting.learned_memory_recommendations_rejected += 1
        else:
            self.accounting.deterministic_fallback_count += 1
        log = {
            "opportunity_id": opportunity_id,
            "recommendation": recommendation.to_dict(),
            "approval": approval.to_dict(),
            "candidate_count": len(candidates),
            "candidate_ids": [candidate.candidate_id for candidate in candidates],
            "executed_by_learned_memory_recommendation": False,
        }
        self.logs.append(log)
        return log

    def summary(self) -> dict[str, Any]:
        accounting = self.accounting.to_dict()
        proof = self.dry_run_proof.to_dict()
        return {
            "enabled": True,
            "dry_run_only": True,
            "recommendation_accounting": accounting,
            "dry_run_proof": proof,
            "learned_memory_writes_executed": accounting["learned_memory_writes_executed"] + proof["learned_memory_writes_executed"],
            "authoritative_memory_write_count": accounting["authoritative_memory_write_count"] + proof["authoritative_memory_write_count"],
            "memory_store_write_count": accounting["memory_store_write_count"] + proof["memory_store_write_count"],
            "file_write_count": accounting["file_write_count"] + proof["file_write_count"],
            "learned_memory_content_generated_count": proof["learned_memory_content_generated_count"],
            "logs": self.logs,
        }


def validate_memory_write_candidate(candidate: MemoryWriteCandidate) -> tuple[bool, str]:
    if not candidate.has_stable_id():
        return False, "candidate_id_not_deterministic"
    if candidate.candidate_type != MEMORY_WRITE_CANDIDATE_TYPE:
        return False, "illegal_candidate_type"
    if candidate.deterministic_generator != DETERMINISTIC_MEMORY_CANDIDATE_GENERATOR:
        return False, "unknown_generator"
    if candidate.memory_type not in SUPPORTED_MEMORY_TYPES:
        return False, "unsupported_memory_type"
    if not candidate.claim.strip():
        return False, "empty_claim"
    if not candidate.dry_run_only:
        return False, "not_dry_run"
    if not candidate.evidence_required:
        return False, "evidence_not_required"
    if not candidate.legal_under_action_mask:
        return False, "action_mask_violation"
    if not candidate.oracle_free:
        return False, "oracle_derived_memory"
    if not candidate.agent_safe:
        return False, "agent_unsafe_memory"
    if not candidate.source_trace_ids or not candidate.source_trace_hashes:
        return False, "missing_source_trace"
    if len(candidate.source_trace_ids) != len(candidate.source_trace_hashes):
        return False, "source_trace_mismatch"
    if not candidate.verifier_evidence:
        return False, "missing_evidence"
    if candidate.confidence < 0.75:
        return False, "low_confidence"
    text = canonical_json(candidate.to_dict()).lower()
    if oracle_leak_keys(candidate.to_dict()):
        return False, "oracle_leak"
    if "canary" in text or "hidden_oracle" in text or "hidden expected" in text:
        return False, "canary_or_oracle_leak"
    try:
        candidate.to_memory_record().validate()
    except ValueError as exc:
        return False, "invalid_memory_record:" + str(exc)
    return True, "ok"


def memory_recommendation_artifact_hash(payload: dict[str, Any]) -> str:
    return stable_hash(payload)


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


def _verifier_evidence_from_traces(traces: list[TraceLedgerRecord]) -> list[dict[str, Any]]:
    evidence: list[dict[str, Any]] = []
    for trace in traces:
        for output in trace.verifier_outputs:
            command = output.get("command")
            success = output.get("success")
            if command is not None:
                evidence.append({"trace_id": trace.trace_id, "command": command, "success": bool(success)})
    if evidence:
        return evidence
    return [{"trace_id": trace.trace_id, "command": ["python", "-m", "pytest", "-q"], "success": trace.success} for trace in traces]


def _dedupe_candidates(candidates: list[MemoryWriteCandidate]) -> list[MemoryWriteCandidate]:
    by_id: dict[str, MemoryWriteCandidate] = {}
    for candidate in candidates:
        by_id.setdefault(candidate.candidate_id, candidate)
    return [by_id[key] for key in sorted(by_id)]
