# memory/scratchpad.py
from pathlib import Path
import json, time

PAD = Path("./memory")
PAD.mkdir(exist_ok=True)

def task_id() -> str:
    return time.strftime("%Y%m%d-%H%M%S")

def write_state(tid: str, state: dict) -> str:
    p = PAD / f"task-{tid}.json"
    p.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
    return str(p)

def read_state(tid: str) -> dict:
    p = PAD / f"task-{tid}.json"
    return json.loads(p.read_text(encoding="utf-8")) if p.exists() else {}

def summarize(state: dict) -> str:
    """把长状态压成给模型看的摘要,控制在几百 token 内。"""
    keep = {k: v for k, v in state.items() if k not in ("raw_docs", "tool_trace")}
    return json.dumps(keep, ensure_ascii=False)[:2000]

tid = task_id()
state = read_state(tid)
state.setdefault("docs", [])
state["docs"].append("docs/report.md")
write_state(tid, state)

# 组装 prompt 时只给摘要,不给全文
prompt = f"当前进度摘要:\n{summarize(state)}\n请继续下一步。"

def chunk_markdown(text: str, max_chars: int = 1200, overlap: int = 150):
    import re
    sections = re.split(r"\n(?=#{1,4}\s)", text)
    chunks, buf = [], ""
    for sec in sections:
        if len(buf) + len(sec) > max_chars and buf:
            chunks.append(buf.strip())
            buf = buf[-overlap:] + "\n" + sec
        else:
            buf = (buf + "\n" + sec) if buf else sec
    if buf.strip():
        chunks.append(buf.strip())
    return chunks

def hybrid_search(query: str, chunks: list[dict], top_k: int = 5, alpha: float = 0.6):
    vec_scores = vector_search(query, chunks)          # 余弦相似度,已归一化到 0~1
    bm25_scores = bm25_search(query, chunks)           # BM25,已归一化到 0~1
    merged = []
    for i, c in enumerate(chunks):
        s = alpha * vec_scores[i] + (1 - alpha) * bm25_scores[i]
        merged.append((s, c))
    merged.sort(key=lambda x: x[0], reverse=True)
    return merged[:top_k]

MEMORY_FIELDS = ["subject", "predicate", "object", "confidence", "source_turn", "expires_at"]

def should_write(fact: dict, existing: list[dict]) -> bool:
    """三条门槛全过才写。"""
    if fact.get("confidence", 0) < 0.8:
        return False
    if fact.get("expires_at") and fact["expires_at"] < today():
        return False
    # 与已有记忆冲突时,不新增,交给 update 流程
    for m in existing:
        if m["subject"] == fact["subject"] and m["predicate"] == fact["predicate"]:
            return False
    return True

def upsert_memory(new_fact: dict, store: list[dict]) -> list[dict]:
    for m in store:
        if m["subject"] == new_fact["subject"] and m["predicate"] == new_fact["predicate"]:
            m["status"] = "superseded"
            m["superseded_at"] = now()
    new_fact["status"] = "active"
    store.append(new_fact)
    return store

import math

def recency_weight(created_at_ts: float, half_life_days: float = 90.0) -> float:
    age_days = (now_ts() - created_at_ts) / 86400
    return 0.5 ** (age_days / half_life_days)

class MemoryAgent:
    def __init__(self, model, store_path="./memory/long_term.json"):
        self.model = model
        self.store_path = store_path
        self.long_term = json.loads(Path(store_path).read_text(encoding="utf-8")) \
            if Path(store_path).exists() else []

    def turn(self, user_input: str, task_id: str):
        # 1) 召回
        recalled = self._recall(user_input, top_k=5)

        # 2) 组装上下文:稳定部分放前面(利于缓存),易变部分放后面
        context = ("## 长期记忆\n" + json.dumps(recalled, ensure_ascii=False)
            + "\n\n## 当前任务状态\n" + summarize(read_state(task_id))
            + "\n\n## 用户输入\n" + user_input
        )

        # 3) 生成
        reply = self.model.generate(context)

        # 4) 回写状态
        state = read_state(task_id)
        state["last_reply"] = reply[:500]
        write_state(task_id, state)

        # 5) 抽取候选记忆 -> 门槛过滤 -> upsert
        candidates = extract_facts(user_input + "\n" + reply)
        for c in candidates:
            if should_write(c, self.long_term):
                self.long_term = upsert_memory(c, self.long_term)
        self._save()
        return reply