from __future__ import annotations

import math
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from typing import Iterable

TOKEN_RE = re.compile(r"[a-z0-9]+")
STOPWORDS = {
    "a", "an", "and", "are", "as", "at", "be", "by", "for", "from",
    "has", "have", "in", "is", "it", "of", "on", "or", "that", "the",
    "this", "to", "with", "you", "your", "will", "when", "if", "can",
    "may", "not", "do", "does", "there", "we", "they", "their", "into",
    "over", "under", "than", "then", "them", "our", "was", "were",
}


def tokenize(text: str) -> list[str]:
    tokens = [t for t in TOKEN_RE.findall(text.lower()) if t not in STOPWORDS]
    return tokens


@dataclass
class SearchHit:
    score: float
    kind: str
    record_id: int
    title: str
    text: str
    citation: str
    metadata: dict


class TfIdfIndex:
    def __init__(self, docs: list[dict[str, object]]):
        self.docs = docs
        self.doc_terms: list[Counter[str]] = []
        self.doc_norms: list[float] = []
        self.idf: dict[str, float] = {}
        self._build()

    def _build(self) -> None:
        doc_freq = defaultdict(int)
        raw_terms: list[Counter[str]] = []
        for doc in self.docs:
            terms = Counter(tokenize(str(doc.get("text", ""))))
            raw_terms.append(terms)
            for term in terms:
                doc_freq[term] += 1
        total_docs = max(len(self.docs), 1)
        self.idf = {
            term: math.log((1 + total_docs) / (1 + df)) + 1.0
            for term, df in doc_freq.items()
        }
        self.doc_terms = raw_terms
        self.doc_norms = [self._norm(tf) for tf in self.doc_terms]

    def _norm(self, tf: Counter[str]) -> float:
        total = 0.0
        for term, count in tf.items():
            total += (count * self.idf.get(term, 0.0)) ** 2
        return math.sqrt(total) or 1.0

    def query(self, text: str, top_k: int = 5) -> list[tuple[int, float]]:
        q_tf = Counter(tokenize(text))
        if not q_tf or not self.docs:
            return []
        q_norm = self._norm(q_tf)
        scores: list[tuple[int, float]] = []
        for i, doc_tf in enumerate(self.doc_terms):
            dot = 0.0
            for term, q_count in q_tf.items():
                if term not in doc_tf:
                    continue
                dot += (q_count * self.idf.get(term, 0.0)) * (doc_tf[term] * self.idf.get(term, 0.0))
            score = dot / (q_norm * self.doc_norms[i])
            if score > 0:
                scores.append((i, score))
        scores.sort(key=lambda item: item[1], reverse=True)
        return scores[:top_k]


class CombinedSearchIndex:
    def __init__(self, section_docs: list[dict[str, object]], job_docs: list[dict[str, object]] | None = None):
        self.section_docs = section_docs
        self.job_docs = job_docs or []
        self.section_index = TfIdfIndex(section_docs)
        self.job_index = TfIdfIndex(job_docs) if job_docs else None

    def search_sections(self, query: str, top_k: int = 5) -> list[SearchHit]:
        hits = []
        for idx, score in self.section_index.query(query, top_k=top_k):
            doc = self.section_docs[idx]
            hits.append(
                SearchHit(
                    score=score,
                    kind="manual_section",
                    record_id=int(doc["record_id"]),
                    title=str(doc["title"]),
                    text=str(doc["text"]),
                    citation=str(doc["citation"]),
                    metadata=dict(doc.get("metadata", {})),
                )
            )
        return hits

    def search_jobs(self, query: str, top_k: int = 5) -> list[SearchHit]:
        if not self.job_index:
            return []
        hits = []
        for idx, score in self.job_index.query(query, top_k=top_k):
            doc = self.job_docs[idx]
            hits.append(
                SearchHit(
                    score=score,
                    kind="job",
                    record_id=int(doc["record_id"]),
                    title=str(doc["title"]),
                    text=str(doc["text"]),
                    citation=str(doc.get("citation", "")),
                    metadata=dict(doc.get("metadata", {})),
                )
            )
        return hits


def merge_search_results(*groups: list[SearchHit], top_k: int = 5) -> list[SearchHit]:
    merged = [hit for group in groups for hit in group]
    merged.sort(key=lambda hit: hit.score, reverse=True)
    return merged[:top_k]