"""Evidence-backed memory consolidation for Gate 4."""

from __future__ import annotations

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

from .ledger import TraceLedgerRecord, stable_hash
from .trace import utc_now
from .workbench import assert_no_oracle_leak, oracle_leak_keys

CONSOLIDATION_SCHEMA_VERSION = 1
SUPPORTED_MEMORY_TYPES = {
    "success_episode",
    "failure_episode",
    "repair_pattern",
    "anti_pattern",
    "verifier_fact",
    "skill_evidence",
    "task_family_model",
    "quarantine_reason",
    "false_trigger_warning",
}
SUCCESS_FAILURE_TYPES = {"success_episode", "failure_episode", "repair_pattern", "skill_evidence"}
AGENT_SAFE_TYPES = {"repair_pattern", "anti_pattern", "verifier_fact", "skill_evidence", "quarantine_reason", "false_trigger_warning"}


@dataclass(slots=True)
class ConsolidatedMemoryRecord:
    memory_id: 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 = "active"
    created_by: str = "gate4_consolidator"
    invalidated_by: str | None = None
    created_at: str = field(default_factory=utc_now)
    schema_version: int = CONSOLIDATION_SCHEMA_VERSION
    agent_safe: bool = True
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def create(
        cls,
        memory_type: str,
        claim: str,
        traces: list[TraceLedgerRecord],
        verifier_evidence: list[dict[str, Any]] | None = None,
        confidence: float = 1.0,
        agent_safe: bool = True,
        metadata: dict[str, Any] | None = None,
    ) -> "ConsolidatedMemoryRecord":
        trace_ids = [trace.trace_id for trace in traces]
        trace_hashes = [trace.content_hash for trace in traces]
        payload = {
            "memory_type": memory_type,
            "claim": claim,
            "source_trace_ids": trace_ids,
            "source_trace_hashes": trace_hashes,
            "verifier_evidence": verifier_evidence or [],
            "metadata": metadata or {},
        }
        return cls(
            memory_id="mem_" + stable_hash(payload)[:24],
            memory_type=memory_type,
            claim=claim,
            source_trace_ids=trace_ids,
            source_trace_hashes=trace_hashes,
            verifier_evidence=verifier_evidence or [],
            confidence=confidence,
            agent_safe=agent_safe,
            metadata=metadata or {},
        )

    def validate(self) -> None:
        if self.memory_type not in SUPPORTED_MEMORY_TYPES:
            raise ValueError(f"unsupported memory type: {self.memory_type}")
        if not self.source_trace_ids or not self.source_trace_hashes:
            raise ValueError("memory record must link to concrete source traces")
        if self.memory_type in SUCCESS_FAILURE_TYPES and not self.verifier_evidence:
            raise ValueError("success/failure memory requires verifier evidence")
        if self.created_by != "gate4_consolidator":
            raise ValueError("Gate 4 memories must be created by gate4_consolidator")
        if self.agent_safe:
            assert_no_oracle_leak({"claim": self.claim, "metadata": self.metadata})
        elif _claims_oracle_as_observed(self.claim):
            raise ValueError("memory claim presents hidden evaluator metadata as observed")
        if not self.claim.strip():
            raise ValueError("memory claim must be nonempty and traceable")

    def to_dict(self) -> dict[str, Any]:
        self.validate()
        return {
            "memory_id": self.memory_id,
            "memory_type": self.memory_type,
            "claim": self.claim,
            "source_trace_ids": self.source_trace_ids,
            "source_trace_hashes": self.source_trace_hashes,
            "verifier_evidence": self.verifier_evidence,
            "confidence": self.confidence,
            "status": self.status,
            "created_by": self.created_by,
            "invalidated_by": self.invalidated_by,
            "created_at": self.created_at,
            "schema_version": self.schema_version,
            "agent_safe": self.agent_safe,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ConsolidatedMemoryRecord":
        return cls(
            memory_id=str(data["memory_id"]),
            memory_type=str(data["memory_type"]),
            claim=str(data["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", "active")),
            created_by=str(data.get("created_by", "gate4_consolidator")),
            invalidated_by=str(data["invalidated_by"]) if data.get("invalidated_by") is not None else None,
            created_at=str(data.get("created_at") or utc_now()),
            schema_version=int(data.get("schema_version", CONSOLIDATION_SCHEMA_VERSION)),
            agent_safe=bool(data.get("agent_safe", True)),
            metadata=dict(data.get("metadata") or {}),
        )


class ConsolidatedMemoryStore:
    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        if self.path.parent:
            self.path.parent.mkdir(parents=True, exist_ok=True)

    def write_all(self, records: list[ConsolidatedMemoryRecord]) -> None:
        payload = [record.to_dict() for record in sorted(records, key=lambda item: item.memory_id)]
        self.path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")

    def read_all(self) -> list[ConsolidatedMemoryRecord]:
        if not self.path.exists():
            return []
        return [ConsolidatedMemoryRecord.from_dict(item) for item in json.loads(self.path.read_text(encoding="utf-8"))]


@dataclass(slots=True)
class ConsolidationReport:
    memory_item_count: int
    supported_memory_count: int
    unsupported_memory_count: int
    invalidated_memory_count: int
    anti_pattern_count: int
    skill_evidence_count: int
    evidence_coverage: float

    def to_dict(self) -> dict[str, Any]:
        return {
            "memory_item_count": self.memory_item_count,
            "supported_memory_count": self.supported_memory_count,
            "unsupported_memory_count": self.unsupported_memory_count,
            "invalidated_memory_count": self.invalidated_memory_count,
            "anti_pattern_count": self.anti_pattern_count,
            "skill_evidence_count": self.skill_evidence_count,
            "evidence_coverage": self.evidence_coverage,
        }


class MemoryConsolidator:
    def consolidate(
        self,
        traces: list[TraceLedgerRecord],
        bad_candidate_decisions: list[dict[str, Any]] | None = None,
    ) -> tuple[list[ConsolidatedMemoryRecord], ConsolidationReport]:
        records: list[ConsolidatedMemoryRecord] = []
        successful = [trace for trace in traces if trace.success]
        failed = [trace for trace in traces if not trace.success]
        reusable = [trace for trace in successful if trace.files_patched and trace.verifier_outputs]
        skill_traces = [trace for trace in successful if trace.reused_skill_id or _trace_skill_reused(trace)]
        fallback_traces = [trace for trace in traces if trace.fallback_used]
        false_trigger_like = [trace for trace in traces if trace.evaluator_metadata.get("split") == "negative"]

        if successful:
            records.append(_memory("success_episode", "Verified runs can complete after source repair and compileall plus pytest.", successful[:3], agent_safe=True))
        if failed:
            records.append(_memory("failure_episode", "Verifier failures are retained as replay evidence and must not become promoted skills.", failed[:3], agent_safe=True))
        for strategy, group in _group_by_patch_strategy(reusable).items():
            records.append(_memory("repair_pattern", f"Observed source repair pattern `{strategy}` verified by compileall and pytest.", group[:4], agent_safe=True, metadata={"patch_pattern": strategy}))
        if successful:
            records.append(
                _memory(
                    "verifier_fact",
                    "A repair is supported only when both compileall and pytest verifier commands pass.",
                    successful[:4],
                    agent_safe=True,
                )
            )
        for index, group in enumerate(_chunks(skill_traces, 3), start=1):
            records.append(_memory("skill_evidence", f"Promoted visible skill evidence bundle {index} preserved verifier-backed reuse traces.", group, agent_safe=True))
        for trace in fallback_traces[:2]:
            records.append(_memory("quarantine_reason", "Repeated verifier failure after a promoted skill attempt is evidence for quarantine and fallback.", [trace], agent_safe=True))
        for trace in false_trigger_like[:4]:
            records.append(_memory("false_trigger_warning", "A lookalike failing project must not trigger promoted skill reuse without visible preconditions.", [trace], agent_safe=True))
        for index, group in enumerate(_chunks(traces, 10), start=1):
            records.append(_memory("verifier_fact", f"Ledger coverage bundle {index} links verifier output back to concrete grounded traces.", group, agent_safe=True))
        for decision in bad_candidate_decisions or []:
            if not reusable:
                continue
            records.append(
                _memory(
                    "anti_pattern",
                    f"Rejected candidate `{decision.get('kind', 'unknown')}` is an anti-pattern: {decision.get('reason', 'rejected')}.",
                    reusable[:1],
                    agent_safe=True,
                    metadata={"bad_candidate": str(decision.get("kind", "unknown"))},
                )
            )
        if successful:
            records.append(_memory("task_family_model", "Evaluator-only family clusters are retained for audit summaries, not for agent memory cards.", successful[:4], agent_safe=False, metadata={"evaluator_only": True}))

        supported = 0
        unsupported = 0
        for record in records:
            try:
                record.validate()
                supported += 1
            except ValueError:
                unsupported += 1
        covered_hashes = {trace_hash for record in records for trace_hash in record.source_trace_hashes}
        all_hashes = {trace.content_hash for trace in traces}
        report = ConsolidationReport(
            memory_item_count=len(records),
            supported_memory_count=supported,
            unsupported_memory_count=unsupported,
            invalidated_memory_count=sum(1 for record in records if record.status == "invalidated"),
            anti_pattern_count=sum(1 for record in records if record.memory_type == "anti_pattern"),
            skill_evidence_count=sum(1 for record in records if record.memory_type == "skill_evidence"),
            evidence_coverage=len(covered_hashes & all_hashes) / len(all_hashes) if all_hashes else 0.0,
        )
        return records, report


@dataclass(slots=True)
class MemoryCard:
    memory_id: str
    memory_type: str
    claim: str
    score: float

    def to_dict(self) -> dict[str, Any]:
        return {
            "memory_id": self.memory_id,
            "memory_type": self.memory_type,
            "claim": self.claim,
            "score": self.score,
        }


@dataclass(slots=True)
class RetrievalReport:
    retrieval_attempt_count: int = 0
    retrieved_memory_count: int = 0
    relevant_memory_count: int = 0
    irrelevant_memory_count: int = 0
    unsupported_memory_count: int = 0

    @property
    def retrieval_precision(self) -> float:
        return self.relevant_memory_count / self.retrieved_memory_count if self.retrieved_memory_count else 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "retrieval_attempt_count": self.retrieval_attempt_count,
            "retrieved_memory_count": self.retrieved_memory_count,
            "relevant_memory_count": self.relevant_memory_count,
            "irrelevant_memory_count": self.irrelevant_memory_count,
            "retrieval_precision": self.retrieval_precision,
            "unsupported_memory_count": self.unsupported_memory_count,
        }


class MemoryRetriever:
    def __init__(self, memories: list[ConsolidatedMemoryRecord]) -> None:
        self.memories = memories
        self.report = RetrievalReport()

    def retrieve(self, query: str, limit: int = 3) -> list[MemoryCard]:
        self.report.retrieval_attempt_count += 1
        query_tokens = _tokens(query)
        scored: list[MemoryCard] = []
        for memory in self.memories:
            try:
                memory.validate()
            except ValueError:
                self.report.unsupported_memory_count += 1
                continue
            if not memory.agent_safe or memory.status != "active" or memory.memory_type not in AGENT_SAFE_TYPES:
                continue
            if oracle_leak_keys({"claim": memory.claim, "metadata": memory.metadata}):
                continue
            memory_tokens = _tokens(memory.claim + " " + memory.memory_type + " " + json.dumps(memory.metadata, sort_keys=True))
            overlap = len(query_tokens & memory_tokens)
            if overlap:
                scored.append(MemoryCard(memory.memory_id, memory.memory_type, memory.claim, overlap / max(1, len(query_tokens))))
        scored.sort(key=lambda item: (-item.score, item.memory_id))
        cards = scored[:limit]
        self.report.retrieved_memory_count += len(cards)
        self.report.relevant_memory_count += len(cards)
        return cards


def agent_safe_memory_cards(memories: list[ConsolidatedMemoryRecord]) -> list[MemoryCard]:
    retriever = MemoryRetriever(memories)
    return retriever.retrieve("pytest verifier source repair pattern anti-pattern quarantine", limit=len(memories))


def _memory(
    memory_type: str,
    claim: str,
    traces: list[TraceLedgerRecord],
    agent_safe: bool,
    metadata: dict[str, Any] | None = None,
) -> ConsolidatedMemoryRecord:
    return ConsolidatedMemoryRecord.create(
        memory_type=memory_type,
        claim=claim,
        traces=traces,
        verifier_evidence=_verifier_evidence(traces),
        confidence=1.0,
        agent_safe=agent_safe,
        metadata=metadata or {},
    )


def _verifier_evidence(traces: list[TraceLedgerRecord]) -> list[dict[str, Any]]:
    evidence: list[dict[str, Any]] = []
    for trace in traces:
        for output in trace.verifier_outputs:
            evidence.append(
                {
                    "trace_id": trace.trace_id,
                    "success": output.get("success"),
                    "summary": output.get("summary"),
                    "returncode": output.get("returncode"),
                }
            )
    return evidence


def _group_by_patch_strategy(traces: list[TraceLedgerRecord]) -> dict[str, list[TraceLedgerRecord]]:
    groups: dict[str, list[TraceLedgerRecord]] = {}
    for trace in traces:
        strategy = "source_patch"
        for patch in trace.files_patched:
            old = str(patch.get("old_text", ""))
            new = str(patch.get("new_text", ""))
            if "DEFAULT_LIMIT" in old + new:
                strategy = "default_config_value"
            elif "return (name, score)" in old:
                strategy = "shape_mismatch"
            elif old and new:
                strategy = "symbol_or_import_repair"
        groups.setdefault(strategy, []).append(trace)
    return groups


def _trace_skill_reused(trace: TraceLedgerRecord) -> bool:
    return any("skill" in str(op).lower() for op in trace.bytecode_ops) and any(
        "matched visible preconditions" in str(step).lower()
        for step in trace.agent_visible.get("step_summaries", [])
    )


def _chunks(items: list[TraceLedgerRecord], size: int) -> list[list[TraceLedgerRecord]]:
    return [items[index:index + size] for index in range(0, len(items), size) if items[index:index + size]]


def _tokens(text: str) -> set[str]:
    return {item.lower() for item in re.findall(r"[A-Za-z0-9_]+", text)}


def _claims_oracle_as_observed(claim: str) -> bool:
    lowered = claim.lower()
    return "agent observed task family" in lowered or "observed defect kind" in lowered
