"""
Text2Mem 实战：实现一个最小可用的记忆操作引擎
运行：python text2mem_practice.py
"""
import json
from datetime import datetime
from typing import ClassVar, Literal
from pydantic import BaseModel, Field, field_validator, model_validator


# ===== 结构与业务校验（生产环境可额外发布 JSON Schema） =====

class MemMeta(BaseModel):
    timestamp: str = ""
    source: str = "user_direct"
    confidence: float = 1.0
    dry_run: bool = False
    confirmation: bool = False

    def model_post_init(self, __context):
        if not self.timestamp:
            self.timestamp = datetime.now().isoformat()


class MemOp(BaseModel):
    """五元组 IR 操作"""
    stage: Literal["ENC", "RET", "STO"]
    op: str
    target: str
    args: dict = Field(default_factory=dict)
    meta: MemMeta = Field(default_factory=MemMeta)

    VALID_OPS: ClassVar[dict[str, set[str]]] = {
        "ENC": {"Encode"},
        "STO": {"Update", "Label", "Promote", "Demote", "Merge", "Split", "Delete", "Lock", "Expire"},
        "RET": {"Retrieve", "Summarize"},
    }

    @field_validator("op")
    @classmethod
    def op_must_be_valid(cls, v):
        all_ops = set()
        for ops in cls.VALID_OPS.values():
            all_ops |= ops
        if v not in all_ops:
            raise ValueError(f"Unknown op: {v}. Valid: {sorted(all_ops)}")
        return v

    @model_validator(mode="after")
    def op_stage_match(self):
        valid = self.VALID_OPS.get(self.stage, set())
        if self.op not in valid:
            raise ValueError(f"Op '{self.op}' not allowed in stage '{self.stage}'. Valid: {sorted(valid)}")
        return self


# ===== 第二层：业务逻辑校验 =====

class MemoryStore:
    """模拟记忆存储引擎；仅实现 Encode/Update/Lock/Delete/Retrieve。"""
    WRITE_OPS: ClassVar[set[str]] = {"Encode", "Update", "Lock", "Delete"}

    def __init__(self):
        self.memories: dict[str, dict] = {}

    def execute(self, ir: MemOp) -> dict:
        if ir.meta.dry_run:
            return {"status": "dry_run", "would_execute": ir.model_dump()}
        if ir.op in self.WRITE_OPS and not ir.meta.confirmation:
            return {"status": "pending_confirmation", "msg": "Write requires external confirmation"}
        handler = getattr(self, f"_exec_{ir.op.lower()}", None)
        if not handler:
            return {"status": "error", "msg": f"No handler for op '{ir.op}'"}
        return handler(ir)

    def _exec_encode(self, ir: MemOp) -> dict:
        mem_id = ir.target or f"mem_{len(self.memories) + 1:04d}"
        self.memories[mem_id] = {
            "id": mem_id, "content": ir.args.get("content", ""),
            "priority": ir.args.get("priority", "normal"),
            "locked": False, "tags": [],
            "created_at": ir.meta.timestamp, "updated_at": ir.meta.timestamp,
        }
        return {"status": "ok", "mem_id": mem_id}

    def _exec_update(self, ir: MemOp) -> dict:
        mem = self._get_or_error(ir.target)
        if isinstance(mem, dict) and "error" in mem:
            return mem
        if mem["locked"]:
            return {"status": "error", "msg": f"Memory '{ir.target}' is locked, cannot update"}
        fields = ir.args.get("fields", {})
        mem.update(fields)
        mem["updated_at"] = ir.meta.timestamp
        return {"status": "ok", "updated": list(fields.keys())}

    def _exec_lock(self, ir: MemOp) -> dict:
        mem = self._get_or_error(ir.target)
        if isinstance(mem, dict) and "error" in mem:
            return mem
        mem["locked"] = True
        return {"status": "ok", "locked": ir.target}

    def _exec_delete(self, ir: MemOp) -> dict:
        mem = self._get_or_error(ir.target)
        if isinstance(mem, dict) and "error" in mem:
            return mem
        if mem["locked"]:
            return {"status": "error", "msg": f"Memory '{ir.target}' is locked, cannot delete"}
        del self.memories[ir.target]
        return {"status": "ok", "deleted": ir.target}

    def _exec_retrieve(self, ir: MemOp) -> dict:
        keyword = ir.args.get("keyword", "")
        results = [m for m in self.memories.values() if keyword.lower() in m["content"].lower()]
        return {"status": "ok", "count": len(results), "results": results}

    def _get_or_error(self, target: str):
        if target not in self.memories:
            return {"status": "error", "msg": f"Memory '{target}' not found"}
        return self.memories[target]


# ===== 运行演示 =====
if __name__ == "__main__":
    store = MemoryStore()
    ops = [
        MemOp(stage="ENC", op="Encode", target="mem_user_lang",
              args={"content": "用户偏好 Python 写后端，TypeScript 写前端", "priority": "high"},
              meta=MemMeta(confirmation=True)),
        MemOp(stage="ENC", op="Encode", target="mem_project",
              args={"content": "当前项目：企业内部智能知识库，使用通义千问和 RAG"},
              meta=MemMeta(confirmation=True)),
        MemOp(stage="STO", op="Lock", target="mem_user_lang",
              args={"reason": "用户身份信息，禁止自动修改"},
              meta=MemMeta(confirmation=True)),
        MemOp(stage="STO", op="Update", target="mem_user_lang",
              args={"fields": {"content": "被篡改的内容"}}),
        MemOp(stage="STO", op="Update", target="mem_user_lang",
              args={"fields": {"content": "仍应被锁拦截"}},
              meta=MemMeta(confirmation=True)),
        MemOp(stage="STO", op="Update", target="mem_project",
              args={"fields": {"content": "当前项目：企业内部智能知识库，检索重排进度 75%"}},
              meta=MemMeta(confirmation=True)),
        MemOp(stage="RET", op="Retrieve", target="", args={"keyword": "知识库"}),
        MemOp(stage="STO", op="Delete", target="mem_project",
              args={}, meta=MemMeta(dry_run=True)),
    ]
    print("=" * 60)
    print("Text2Mem 实战：记忆操作引擎演示")
    print("=" * 60)
    for i, op in enumerate(ops, 1):
        print(f"\n--- 操作 {i}: {op.stage}/{op.op} ---")
        result = store.execute(op)
        print(f"  结果: {json.dumps(result, ensure_ascii=False, indent=2)}")
    print(f"\n最终记忆库状态（{len(store.memories)} 条）：")
    for mid, mem in store.memories.items():
        locked = " 🔒" if mem["locked"] else ""
        print(f"  [{mid}]{locked}: {mem['content']}")