# retoor <retoor@molodetz.nl>

from __future__ import annotations

import collections
import logging
import math
import re
import uuid
from typing import Any

from ..tasks.schedule import now_utc, to_iso

logger = logging.getLogger("devii.agentic.lessons")

TABLE = "devii_lessons"
TOKEN_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|\d+")
BM25_K1 = 1.5
BM25_B = 0.75


def tokenize(text: str) -> list[str]:
    tokens = TOKEN_RE.findall((text or "").lower())
    extra: list[str] = []
    for token in tokens:
        extra.extend(re.findall(r"[a-z]+", token))
    return list(dict.fromkeys(tokens + extra))


class LessonStore:
    def __init__(self, db: Any, owner_kind: str, owner_id: str) -> None:
        self._db = db
        self._owner_kind = owner_kind
        self._owner_id = owner_id
        self._dirty = True
        self._docs: list[dict[str, Any]] = []
        self._tf: list[collections.Counter] = []
        self._dl: list[int] = []
        self._idf: dict[str, float] = {}
        self._avgdl = 0.0
        self._n = 0
        self._ensure_indexes()

    def _ensure_indexes(self) -> None:
        if TABLE not in self._db.tables:
            return
        table = self._db[TABLE]
        if not table.has_column("deleted_at"):
            table.create_column_by_example("deleted_at", "")
        if not table.has_column("deleted_by"):
            table.create_column_by_example("deleted_by", "")
        table.create_index(["owner_kind", "owner_id"])

    @property
    def _table(self) -> Any:
        return self._db[TABLE]

    @property
    def _scope(self) -> dict[str, str]:
        return {"owner_kind": self._owner_kind, "owner_id": self._owner_id}

    def count(self) -> int:
        if TABLE not in self._db.tables:
            return 0
        return self._table.count(deleted_at=None, **self._scope)

    def add(
        self, observation: str, conclusion: str, next_action: str, tags: str = ""
    ) -> dict[str, Any]:
        record = {
            "uid": uuid.uuid4().hex,
            "observation": observation,
            "conclusion": conclusion,
            "next_action": next_action,
            "tags": tags,
            "created_at": to_iso(now_utc()),
            "hits": 0,
            "deleted_at": None,
            "deleted_by": None,
            **self._scope,
        }
        self._table.insert(record)
        self._dirty = True
        logger.info(
            "Lesson stored owner=%s/%s tags=%s", self._owner_kind, self._owner_id, tags
        )
        return record

    def all(self) -> list[dict[str, Any]]:
        if TABLE not in self._db.tables:
            return []
        return list(self._table.find(deleted_at=None, **self._scope))

    def delete(self, uid: str) -> bool:
        if TABLE not in self._db.tables:
            return False
        row = self._table.find_one(uid=uid, deleted_at=None, **self._scope)
        if not row:
            return False
        self._table.update(
            {
                "id": row["id"],
                "deleted_at": to_iso(now_utc()),
                "deleted_by": f"{self._owner_kind}:{self._owner_id}",
            },
            ["id"],
        )
        self._dirty = True
        return True

    def clear(self) -> int:
        n = self.count()
        if TABLE in self._db.tables:
            self._table.delete(**self._scope)
        self._dirty = True
        logger.info(
            "Cleared %d lesson(s) for owner=%s/%s", n, self._owner_kind, self._owner_id
        )
        return n

    def _rebuild(self) -> None:
        rows = self.all()
        df: collections.Counter = collections.Counter()
        docs: list[dict[str, Any]] = []
        tf_list: list[collections.Counter] = []
        dl: list[int] = []
        for row in rows:
            text = " ".join(
                str(row.get(field) or "")
                for field in ("observation", "conclusion", "next_action", "tags")
            )
            tokens = tokenize(text)
            if not tokens:
                continue
            tf = collections.Counter(tokens)
            for term in tf:
                df[term] += 1
            docs.append(row)
            tf_list.append(tf)
            dl.append(len(tokens))
        self._docs = docs
        self._tf = tf_list
        self._dl = dl
        self._n = len(docs)
        self._avgdl = sum(dl) / max(self._n, 1)
        self._idf = {
            term: math.log((self._n - freq + 0.5) / (freq + 0.5) + 1)
            for term, freq in df.items()
        }
        self._dirty = False

    def search(self, query: str, k: int = 3) -> list[dict[str, Any]]:
        if self._dirty:
            self._rebuild()
        tokens = tokenize(query)
        if not tokens or not self._docs:
            return []
        scored: list[tuple[float, int]] = []
        for index, tf in enumerate(self._tf):
            score = 0.0
            for token in tokens:
                freq = tf.get(token, 0)
                if freq == 0:
                    continue
                idf = self._idf.get(token, 0.0)
                norm = 1 - BM25_B + BM25_B * (self._dl[index] / max(self._avgdl, 1))
                score += idf * (freq * (BM25_K1 + 1)) / (freq + BM25_K1 * norm)
            if score > 0:
                scored.append((score, index))
        scored.sort(reverse=True)
        results: list[dict[str, Any]] = []
        for score, index in scored[:k]:
            row = self._docs[index]
            results.append(
                {
                    "uid": row.get("uid"),
                    "observation": row.get("observation"),
                    "conclusion": row.get("conclusion"),
                    "next_action": row.get("next_action"),
                    "tags": row.get("tags"),
                    "score": round(score, 3),
                }
            )
        return results
