"""Base contract for memory methods.

A method implements memorize() for each visit and build_memory_context() for each question.
The runner owns the answering call, so every method is answered with the same prompt, except
methods whose own agent produces the answer (ANSWERS_INTERNALLY). memorize() receives the
rendered visit text with its visit number and date.
"""
from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field


def as_list(x) -> list:
    """Coerce anything list-like (a numpy array included) to a plain list."""
    if x is None:
        return []
    try:
        return list(x)
    except Exception:
        return []


class Tokenizer:
    """Token counter backed by the tiktoken encoding named in backbone.yaml (`answer.tokenizer` for context budgets
    and chunk sizes, `embedding.tokenizer` for embedding inputs); run_eval sets it before any method runs."""

    def __init__(self, role: str):
        self.role, self._enc = role, None

    def set(self, name: str | None) -> None:
        if name and not (isinstance(name, str) and name.startswith("<")):
            import tiktoken
            self._enc = tiktoken.get_encoding(name)
        else:
            raise SystemExit(f"set {self.role}.tokenizer in backbone.yaml")

    def _get(self):
        if self._enc is None:
            raise RuntimeError(f"{self.role}.tokenizer is not set (backbone.yaml)")
        return self._enc

    def encode(self, text: str, **kw) -> list[int]:
        return self._get().encode(text, **kw)

    def decode(self, tokens) -> str:
        return self._get().decode(tokens)


ENC = Tokenizer("answer")
EMBED_ENC = Tokenizer("embedding")


@dataclass
class MemoryBuildResult:
    """Result of one memorize() call."""
    success: bool = True
    method: str = ""
    action: str = ""
    visit_number: int = -1
    chunk_count: int = 0
    added_chunks: int = 0
    time_cost: float = 0.0
    extra: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {"success": self.success, "method": self.method, "action": self.action,
                "visit_number": self.visit_number, "chunk_count": self.chunk_count,
                "added_chunks": self.added_chunks, "time_cost": round(self.time_cost, 4),
                **self.extra}


@dataclass
class MemoryContext:
    """The MEMORY block handed to the runner, or the answer of an internally answering method."""
    mem: str
    n_tok: int = 0
    retrieved_count: int = 0
    retrieved_visits: list = field(default_factory=list)
    retrieved_memories: list = field(default_factory=list)
    context_truncated: bool = False
    error: bool = False
    answer: str | None = None
    extra: dict = field(default_factory=dict)


class BaseAgent(ABC):
    ANSWERS_INTERNALLY = False
    MEMORIZE_WITH_LLM = False

    def __init__(self, **params):
        self.params = params
        self._context_id: str | None = None

    def set_trace_sink(self, sink) -> None:
        """Install the callback that appends one event per dict to `memory_trace.jsonl`."""
        self._trace_sink = sink

    def _trace(self, event: str, **fields) -> None:
        sink = getattr(self, "_trace_sink", None)
        if sink is None:
            return
        try:
            from ..llm_client import USAGE
            sink({"method": getattr(self, "METHOD_NAME", type(self).__name__),
                  "event": event, "phase": getattr(USAGE, "phase", None), **fields})
        except Exception as e:
            n = getattr(self, "_trace_fail", 0) + 1
            self._trace_fail = n
            if n == 1:
                try:
                    sink({"method": getattr(self, "METHOD_NAME", type(self).__name__),
                          "event": "trace_error", "error": f"{type(e).__name__}: {e}"[:300],
                          "failed_event": event})
                except Exception:
                    pass

    def consolidate(self) -> dict | None:
        """Sleep-time memory maintenance, run at each checkpoint before answering."""
        return None

    def dump_memory(self, out_dir) -> dict | None:
        """Export the current memory state into out_dir; return a small stats dict, or None."""
        return None

    def set_context_id(self, context_id: str) -> None:
        """Set the per-patient namespace."""
        self._context_id = str(context_id)

    @abstractmethod
    def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult:
        """Store one rendered visit into memory."""

    def memorize_session(self, session: dict,
                         rendered_text: str | None = None) -> MemoryBuildResult:
        """Memorize one session from dataset.load_sessions()."""
        if rendered_text is None:
            from ..dataset import render_visit
            rendered_text = render_visit(session)
        return self.memorize(rendered_text, session["visit_number"], session["date"])

    @abstractmethod
    def build_memory_context(self, question: str, as_of: str = "",
                             display_question: str | None = None) -> MemoryContext:
        """Produce the MEMORY block for one question."""

    @staticmethod
    def count_tokens(text: str) -> int:
        return len(ENC.encode(text)) if text else 0

    @staticmethod
    def cut_tokens(text: str, max_tokens: int) -> str:
        """The first `max_tokens` tokens of `text`."""
        toks = ENC.encode(text)
        return text if len(toks) <= max_tokens else ENC.decode(toks[:max_tokens])

    @staticmethod
    def split_tokens(text: str, chunk_size: int, chunk_overlap: int,
                     header: str = "") -> list[str]:
        """Split text into token windows; continuation chunks repeat `header` + '(cont.)'."""
        if chunk_overlap >= chunk_size:
            raise ValueError(f"chunk_overlap ({chunk_overlap}) must be smaller "
                             f"than chunk_size ({chunk_size})")
        if not text.strip():
            return []
        tokens = ENC.encode(text)
        if len(tokens) <= chunk_size:
            return [text]
        chunks, start = [], 0
        while start < len(tokens):
            end = min(start + chunk_size, len(tokens))
            piece = ENC.decode(tokens[start:end])
            if start > 0 and header:
                piece = f"{header} (cont.)\n{piece}"
            chunks.append(piece)
            if end == len(tokens):
                break
            start = end - chunk_overlap
        return chunks


def snapshot_jsonl_gz(path, rows) -> int:
    """Write rows as gzipped JSONL; returns row count."""
    import gzip
    import json as _json
    with gzip.open(path, "wt", encoding="utf-8") as fh:
        n = 0
        for r in rows:
            fh.write(_json.dumps(r, ensure_ascii=False, default=str) + "\n")
            n += 1
    return n


def snapshot_sqlite(src, dst) -> None:
    """Copy an open SQLite database with the backup API."""
    import sqlite3
    a = sqlite3.connect(f"file:{src}?mode=ro", uri=True)
    try:
        b = sqlite3.connect(str(dst))
        try:
            with b:
                a.backup(b)
        finally:
            b.close()
    finally:
        a.close()


def snapshot_vectors_delta(out_dir, pairs, seen: set) -> int:
    """Write the vectors this checkpoint added, and record them as seen."""
    import numpy as np
    ids, vecs = [], []
    for mid, v in pairs:
        mid = str(mid)
        if not mid or mid in seen or v is None:
            continue
        seen.add(mid)
        ids.append(mid)
        vecs.append(np.asarray(v, dtype="float32"))
    if not ids:
        np.savez_compressed(out_dir / "vectors.npz",
                            ids=np.array([], dtype=object), vecs=np.zeros((0, 0), dtype="float32"))
        return 0
    np.savez_compressed(out_dir / "vectors.npz",
                        ids=np.array(ids, dtype=object), vecs=np.vstack(vecs))
    return len(ids)


def snapshot_tree(src, dst) -> int:
    """Copy a store directory, SQLite files through the backup API, skipping lock and tmp files; return the count."""
    import shutil
    from pathlib import Path as _P
    src, dst = _P(src), _P(dst)
    n = 0
    for f in src.rglob("*"):
        if not f.is_file() or f.name.endswith((".lock", "-wal", "-shm")):
            continue
        if "tmp" in f.relative_to(src).parts[:-1]:
            continue
        rel = dst / f.relative_to(src)
        rel.parent.mkdir(parents=True, exist_ok=True)
        if f.suffix in (".db", ".sqlite") or f.name.endswith(".sqlite3") or f.name == "storage.sqlite":
            try:
                snapshot_sqlite(f, rel)
            except Exception:
                shutil.copy2(f, rel)
        else:
            shutil.copy2(f, rel)
        n += 1
    return n