# retoor <retoor@molodetz.nl>

from __future__ import annotations

import collections
import logging
import math
import re
import uuid
from datetime import timedelta
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

DEDUP_JACCARD_THRESHOLD = 0.70
DEFAULT_MAX_PER_OWNER = 500
DEFAULT_MAX_AGE_DAYS = 90
LOW_QUALITY_THRESHOLD = -3


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))


def _jaccard(tokens_a: set[str], tokens_b: set[str]) -> float:
    if not tokens_a and not tokens_b:
        return 0.0
    if not tokens_a or not tokens_b:
        return 0.0
    return len(tokens_a & tokens_b) / len(tokens_a | tokens_b)


def _read_retention_settings(db: Any) -> tuple[int, int]:
    max_per = DEFAULT_MAX_PER_OWNER
    max_age = DEFAULT_MAX_AGE_DAYS
    if "site_settings" not in db.tables:
        return max_per, max_age
    for row in db["site_settings"].find(key={"in": ["devii_lessons_max_per_owner", "devii_lessons_max_age_days"]}):
        if row["key"] == "devii_lessons_max_per_owner":
            try:
                max_per = int(row["value"])
            except (ValueError, TypeError):
                pass
        elif row["key"] == "devii_lessons_max_age_days":
            try:
                max_age = int(row["value"])
            except (ValueError, TypeError):
                pass
    return max_per, max_age


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_columns()
        self._ensure_indexes()

    def _ensure_columns(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", "")
        if not table.has_column("rating"):
            table.create_column_by_example("rating", 0)

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

    @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 _row_text(self, row: dict[str, Any]) -> str:
        return " ".join(
            str(row.get(field) or "")
            for field in ("observation", "conclusion", "next_action", "tags")
        )

    def _find_similar(self, text: str, threshold: float = DEDUP_JACCARD_THRESHOLD) -> dict[str, Any] | None:
        query_tokens = set(tokenize(text))
        if not query_tokens:
            return None
        all_rows = self.all()
        best_row: dict[str, Any] | None = None
        best_score = 0.0
        for row in all_rows:
            row_tokens = set(tokenize(self._row_text(row)))
            score = _jaccard(query_tokens, row_tokens)
            if score > best_score and score >= threshold:
                best_score = score
                best_row = row
        return best_row

    def _enforce_cap(self, max_per_owner: int) -> int:
        soft_deleted = 0
        while True:
            current = self.count()
            if current <= max_per_owner:
                break
            excess = current - max_per_owner
            rows = list(
                self._table.find(
                    deleted_at=None,
                    order_by=["created_at"],
                    _limit=excess,
                    **self._scope,
                )
            )
            if not rows:
                break
            now = to_iso(now_utc())
            for row in rows:
                self._table.update(
                    {
                        "id": row["id"],
                        "deleted_at": now,
                        "deleted_by": "retention",
                    },
                    ["id"],
                )
                soft_deleted += 1
            self._dirty = True
        if soft_deleted:
            logger.info(
                "Retention cap pruned %d lesson(s) for owner=%s/%s",
                soft_deleted,
                self._owner_kind,
                self._owner_id,
            )
        return soft_deleted

    def add(
        self, observation: str, conclusion: str, next_action: str, tags: str = ""
    ) -> dict[str, Any]:
        text = " ".join([observation, conclusion, next_action, tags])
        similar = self._find_similar(text)
        if similar and similar.get("id") is not None:
            hits = (similar.get("hits") or 0) + 1
            self._table.update(
                {
                    "id": similar["id"],
                    "hits": hits,
                    "created_at": to_iso(now_utc()),
                },
                ["id"],
            )
            self._dirty = True
            logger.info(
                "Lesson deduplicated owner=%s/%s hits=%d",
                self._owner_kind,
                self._owner_id,
                hits,
            )
            return {**similar, "hits": hits, "deduplicated": True}

        uid = uuid.uuid4().hex
        record = {
            "uid": uid,
            "observation": observation,
            "conclusion": conclusion,
            "next_action": next_action,
            "tags": tags,
            "created_at": to_iso(now_utc()),
            "hits": 0,
            "rating": 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
        )
        max_per, _ = _read_retention_settings(self._db)
        self._enforce_cap(max_per)
        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 rate(self, uid: str, value: int) -> 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
        current = row.get("rating") or 0
        self._table.update(
            {"id": row["id"], "rating": current + value},
            ["id"],
        )
        self._dirty = True
        logger.info(
            "Lesson %s rated %+d (now %d) owner=%s/%s",
            uid,
            value,
            current + value,
            self._owner_kind,
            self._owner_id,
        )
        return True

    def prune(self, max_age_days: int | None = None) -> int:
        if TABLE not in self._db.tables:
            return 0
        if max_age_days is None:
            _, max_age_days = _read_retention_settings(self._db)
        cutoff = now_utc() - timedelta(days=max_age_days)
        cutoff_iso = to_iso(cutoff)
        rows = list(
            self._table.find(
                deleted_at=None,
                created_at={"<": cutoff_iso},
                **self._scope,
            )
        )
        if not rows:
            return 0
        now = to_iso(now_utc())
        soft_deleted = 0
        for row in rows:
            self._table.update(
                {
                    "id": row["id"],
                    "deleted_at": now,
                    "deleted_by": "retention",
                },
                ["id"],
            )
            soft_deleted += 1
        if soft_deleted:
            self._dirty = True
            logger.info(
                "Pruned %d old lesson(s) for owner=%s/%s",
                soft_deleted,
                self._owner_kind,
                self._owner_id,
            )
        return soft_deleted

    def prune_all_owners(self, max_age_days: int | None = None) -> int:
        if TABLE not in self._db.tables:
            return 0
        if max_age_days is None:
            _, max_age_days = _read_retention_settings(self._db)
        cutoff = now_utc() - timedelta(days=max_age_days)
        cutoff_iso = to_iso(cutoff)
        rows = list(
            self._table.find(
                deleted_at=None,
                created_at={"<": cutoff_iso},
            )
        )
        if not rows:
            return 0
        now = to_iso(now_utc())
        soft_deleted = 0
        for row in rows:
            self._table.update(
                {
                    "id": row["id"],
                    "deleted_at": now,
                    "deleted_by": "retention",
                },
                ["id"],
            )
            soft_deleted += 1
        if soft_deleted:
            logger.info("Pruned %d old lesson(s) across all owners", soft_deleted)
        return soft_deleted

    def _rebuild(self) -> None:
        rows = self.all()
        df: collections.Counter = collections.Counter()
        docs: list[dict[str, Any]] = []
        tf_list: list[collections.Counter] = []
        dl_list: list[int] = []
        for row in rows:
            rating = row.get("rating") or 0
            if rating <= LOW_QUALITY_THRESHOLD:
                continue
            text = self._row_text(row)
            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_list.append(len(tokens))
        self._docs = docs
        self._tf = tf_list
        self._dl = dl_list
        self._n = len(docs)
        self._avgdl = sum(dl_list) / 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
