from __future__ import annotations

import math
import re

import numpy as np

from .extract import extract_module
from .models import Graph, Module, Topic

_NAME_STOP = {"course", "material", "slide", "slides", "syllabus", "pdf", "pptx", "docx", "doc"}
_NAME_WEIGHT = 2.0   # a clearly-matching filename must beat any plausible content-cosine disagreement


def _slug(s: str) -> str:
    out = re.sub(r"[^a-z0-9]+", "-", s.lower()).strip("-")
    return out[:60] or "x"


def _tokens(s: str) -> set[str]:
    s = re.sub(r"(?<=[a-z])(?=[A-Z])", " ", s)          # MarcoGiordani -> Marco Giordani
    toks = re.findall(r"[a-z0-9]+", s.lower())
    return {t for t in toks if t not in _NAME_STOP and (len(t) > 1 or t.isdigit())}


def _name_scores(file_path: str, identities: list[set[str]], weights: dict[str, float]) -> np.ndarray:
    """How much of each module's identity (title+lecturer+syllabus filename tokens, IDF-weighted)
    is covered by the file's path tokens. Rare tokens — lecturer names, '1)'/'2)' prefixes —
    dominate; ubiquitous ones ('5g', 'network') barely count. 0 when nothing matches."""
    ft = _tokens(file_path)
    out = np.zeros(len(identities), dtype="float32")
    for i, ident in enumerate(identities):
        total = sum(weights[t] for t in ident)
        if total > 0:
            out[i] = sum(weights[t] for t in ident & ft) / total
    return out


def build_graph(chunks, chunk_vectors, syllabi: dict, embedder, llm) -> Graph:
    """Build the Module->Topic->Source->Chunk graph.

    chunks: list[Chunk]; chunk_vectors: array aligned with chunks (reused index embeddings);
    syllabi: {syllabus_file_path: full_text}. Modules + topics come from the syllabi (LLM).
    Each content file is matched to one module by filename affinity + content similarity —
    name affinity matters because sibling modules of one course often have near-identical
    titles, where content embeddings alone mis-route whole files. Each chunk then goes to
    its nearest topic within that module (top-1, never left unlinked)."""
    modules: list[Module] = []
    topics: list[Topic] = []
    topic_titles: list[str] = []

    for path, text in syllabi.items():
        spec = extract_module(text, llm)
        mid = _slug(path)
        module = Module(id=mid, title=spec.title or path, lecturer=spec.lecturer, hours=spec.hours,
                        objective=spec.objective, language=spec.language, source_file=path)
        for title in spec.topics:
            tid = f"{mid}::{_slug(title)}"
            topics.append(Topic(id=tid, module_id=mid, title=title))
            topic_titles.append(f"{spec.title}. {title}")
            module.topic_ids.append(tid)
        modules.append(module)

    if not modules or not topics:
        return Graph(modules=modules, topics=topics)

    topic_vecs = embedder.encode(topic_titles).astype("float32")
    module_vecs = embedder.encode([f"{m.title}. {m.lecturer}. {m.objective}" for m in modules]).astype("float32")
    module_vecs /= np.linalg.norm(module_vecs, axis=1, keepdims=True) + 1e-9

    identities = [_tokens(f"{m.title} {m.lecturer} {m.source_file}") for m in modules]
    df: dict[str, int] = {}
    for ident in identities:
        for t in ident:
            df[t] = df.get(t, 0) + 1
    idf = {t: math.log(1 + len(modules) / n) for t, n in df.items()}

    syllabus_files = set(syllabi.keys())
    by_file: dict[str, list] = {}
    for chunk, vec in zip(chunks, chunk_vectors):
        if chunk.file in syllabus_files:
            continue
        by_file.setdefault(chunk.file, []).append((chunk, np.asarray(vec, dtype="float32")))

    topic_module = [t.module_id for t in topics]
    topic_by_id = {t.id: t for t in topics}

    for file, items in by_file.items():
        file_vec = np.mean([v for _, v in items], axis=0)
        file_vec /= np.linalg.norm(file_vec) + 1e-9
        score = _NAME_WEIGHT * _name_scores(file, identities, idf) + module_vecs @ file_vec
        module = modules[int(np.argmax(score))]
        module.source_ids.append(file)
        local = [i for i, m in enumerate(topic_module) if m == module.id]
        if not local:
            continue
        sub = topic_vecs[local]
        for chunk, vec in items:
            j = int(np.argmax(sub @ vec))   # nearest topic in this module — never leaves a chunk unlinked
            topic_by_id[topics[local[j]].id].chunk_ids.append(chunk.id)

    return Graph(modules=modules, topics=topics)