"""TF-IDF vector log search index for otelctx."""

from __future__ import annotations

import math
import re
import threading
from collections import defaultdict
from dataclasses import dataclass, field


@dataclass
class IndexedLog:
    log_id: int
    body: str
    service_name: str
    severity: str
    timestamp: str
    attributes: dict[str, str] = field(default_factory=dict)
    _tokens: list[str] = field(default_factory=list, repr=False)


def _tokenize(text: str) -> list[str]:
    text = text.lower()
    return re.findall(r"[a-z0-9_.\-/]+", text)


class LogVectorIndex:
    """In-memory TF-IDF index for log full-text search with FIFO eviction."""

    def __init__(self, capacity: int = 50_000) -> None:
        self.capacity = capacity
        self._lock = threading.RLock()
        self._docs: list[IndexedLog] = []
        self._log_id_set: set[int] = set()
        # doc_freq[token] = number of docs containing token
        self._doc_freq: dict[str, int] = defaultdict(int)
        # tf[doc_idx][token] = term frequency in doc
        self._tf: list[dict[str, float]] = []

    def _compute_tf(self, tokens: list[str]) -> dict[str, float]:
        if not tokens:
            return {}
        freq: dict[str, int] = defaultdict(int)
        for t in tokens:
            freq[t] += 1
        n = len(tokens)
        return {t: c / n for t, c in freq.items()}

    def _update_df(self, tf: dict[str, float], delta: int = 1) -> None:
        for token in tf:
            self._doc_freq[token] += delta
            if self._doc_freq[token] <= 0:
                del self._doc_freq[token]

    def add(self, log: IndexedLog) -> None:
        with self._lock:
            if log.log_id in self._log_id_set:
                return
            # Evict oldest if at capacity
            if len(self._docs) >= self.capacity:
                evicted = self._docs.pop(0)
                self._log_id_set.discard(evicted.log_id)
                old_tf = self._tf.pop(0)
                self._update_df(old_tf, delta=-1)

            tokens = _tokenize(log.body)
            log._tokens = tokens
            tf = self._compute_tf(tokens)
            self._update_df(tf, delta=1)
            self._docs.append(log)
            self._log_id_set.add(log.log_id)
            self._tf.append(tf)

    def add_many(self, logs: list[IndexedLog]) -> int:
        count = 0
        for log in logs:
            self.add(log)
            count += 1
        return count

    def _idf(self, token: str) -> float:
        n = len(self._docs)
        df = self._doc_freq.get(token, 0)
        if df == 0 or n == 0:
            return 0.0
        return math.log((n + 1) / (df + 1)) + 1.0

    def _tfidf_vec(self, tf: dict[str, float]) -> dict[str, float]:
        return {t: v * self._idf(t) for t, v in tf.items()}

    def _cosine(self, a: dict[str, float], b: dict[str, float]) -> float:
        common = set(a) & set(b)
        if not common:
            return 0.0
        dot = sum(a[t] * b[t] for t in common)
        mag_a = math.sqrt(sum(v * v for v in a.values()))
        mag_b = math.sqrt(sum(v * v for v in b.values()))
        if mag_a == 0 or mag_b == 0:
            return 0.0
        return dot / (mag_a * mag_b)

    def search(
        self,
        query: str,
        top_k: int = 10,
        service_name: str | None = None,
        severity: str | None = None,
        min_score: float = 0.01,
    ) -> list[dict]:
        with self._lock:
            qtokens = _tokenize(query)
            if not qtokens:
                return []
            qtf = self._compute_tf(qtokens)
            q_vec = self._tfidf_vec(qtf)

            results = []
            for i, doc in enumerate(self._docs):
                if service_name and doc.service_name != service_name:
                    continue
                if severity and doc.severity.upper() != severity.upper():
                    continue
                doc_vec = self._tfidf_vec(self._tf[i])
                score = self._cosine(q_vec, doc_vec)
                if score >= min_score:
                    results.append((score, doc))

            results.sort(key=lambda x: x[0], reverse=True)
            return [
                {
                    "score": round(s, 4),
                    "log_id": d.log_id,
                    "body": d.body,
                    "service_name": d.service_name,
                    "severity": d.severity,
                    "timestamp": d.timestamp,
                }
                for s, d in results[:top_k]
            ]

    def size(self) -> int:
        with self._lock:
            return len(self._docs)
