"""Lecture-JSONL chunker.

Reads upstream `lecture_memory_*.jsonl` files (session_header + turn +
problem_outcome records) and emits per-problem retrieval citizens and
per-turn citation targets, plus a (session_date, row_index) -> doc_ids
lookup map for syntopicon reference resolution.
"""
from __future__ import annotations
import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path


@dataclass
class ChunkerOutput:
    lecture_problems: list[dict] = field(default_factory=list)
    lecture_turns: list[dict] = field(default_factory=list)
    # (session_date, row_index) -> (lecture_turn_doc_id, lecture_problem_doc_id)
    row_lookup: dict[tuple[str, int], tuple[str, str]] = field(default_factory=dict)


def _hash8(*parts: str) -> str:
    h = hashlib.sha256("|".join(parts).encode("utf-8"))
    return h.hexdigest()[:8]


def _problem_doc_id(session_date: str, problem_id: str) -> str:
    return f"lecture_problem_{session_date}_{problem_id}_{_hash8(session_date, problem_id)}"


def _turn_doc_id(session_date: str, row_index: int) -> str:
    return f"lecture_turn_{session_date}_R{row_index:04d}_{_hash8(session_date, str(row_index))}"


def chunk_lecture_jsonl(
    path: Path,
    *,
    term_registry: set[str] | None = None,
) -> ChunkerOutput:
    """Read a lecture_memory_*.jsonl file and produce per-problem + per-turn records.

    term_registry: if provided, lecture_problem.concepts_present is filtered
    to terms in this set. If None, concepts_present is the raw union of upstream
    'terms' fields (used by tests that don't care about registry filtering).
    """
    out = ChunkerOutput()
    sessions: dict[str, dict] = {}  # session_date -> session_header dict
    # (session_date, problem_id) -> list of turn dicts
    problems: dict[tuple[str, str], list[dict]] = {}
    # session_date -> instructor_name (for fallback if turn lacks)
    session_instructor: dict[str, str] = {}

    for line_no, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        raw = raw.strip()
        if not raw:
            continue
        try:
            rec = json.loads(raw)
        except json.JSONDecodeError as e:
            raise ValueError(
                f"malformed JSONL at {path}:{line_no}: {e.msg}"
            ) from e

        rt = rec.get("record_type")
        if rt == "session_header":
            sd = rec.get("lecture_date")
            sessions[sd] = rec
            session_instructor[sd] = rec.get("instructor_name", "")
        elif rt == "turn":
            sd = rec.get("lecture_date")
            pid = rec.get("problem_id") or "P_unstructured"
            problems.setdefault((sd, pid), []).append(rec)
        # problem_outcome records ignored (no retrieval value at chunk level)

    # Emit lecture_turn records and lecture_problem records.
    for (session_date, problem_id), turns in problems.items():
        if not turns:
            continue
        # Sort by row_index ascending
        turns_sorted = sorted(turns, key=lambda t: t.get("row_index", 0))
        problem_doc_id = _problem_doc_id(session_date, problem_id)
        child_turn_doc_ids: list[str] = []
        for t in turns_sorted:
            if "row_index" not in t:
                raise ValueError(
                    f"turn record in {path} session {session_date} "
                    f"problem {problem_id} is missing required field 'row_index'"
                )
            row_index = t["row_index"]
            turn_doc_id = _turn_doc_id(session_date, row_index)
            child_turn_doc_ids.append(turn_doc_id)
            out.lecture_turns.append({
                "kind": "lecture_turn",
                "doc_id": turn_doc_id,
                "lecture_date": session_date,
                "row_index": row_index,
                "turn_id": t.get("turn_id", ""),
                "start": t.get("start"),
                "end": t.get("end"),
                "speaker": t.get("speaker", ""),
                "label": t.get("label", ""),
                "text": t.get("text", ""),
                "word_count": t.get("word_count", 0),
                "duration": t.get("duration"),
                "parent_problem_doc_id": problem_doc_id,
            })
            out.row_lookup[(session_date, row_index)] = (turn_doc_id, problem_doc_id)

        text_body = "\n".join(t.get("text", "") for t in turns_sorted)
        terms_union: set[str] = set()
        for t in turns_sorted:
            for term in t.get("terms", []) or []:
                # The upstream JSONL may emit terms as plain strings (test fixtures)
                # or as dicts with a "term" key (production data schema v1.0+).
                if isinstance(term, dict):
                    slug = term.get("term", "")
                    if slug:
                        terms_union.add(slug)
                elif isinstance(term, str) and term:
                    terms_union.add(term)
        if term_registry is not None:
            terms_present = sorted(terms_union & term_registry)
        else:
            terms_present = sorted(terms_union)
        first_q = next(
            (t.get("problem_question", "") for t in turns_sorted if t.get("problem_question")),
            "",
        )
        session = sessions.get(session_date, {})
        row_indices = [t["row_index"] for t in turns_sorted]
        out.lecture_problems.append({
            "kind": "lecture_problem",
            "doc_id": problem_doc_id,
            "lecture_date": session_date,
            "lecture_index_of_total": session.get("lecture_index_of_total", ""),
            "problem_id": problem_id,
            "problem_question": first_q,
            "instructor_name": session_instructor.get(session_date, ""),
            "child_turn_doc_ids": child_turn_doc_ids,
            "row_range": [min(row_indices), max(row_indices)],
            "text": text_body,
            "concepts_present": terms_present,
        })

    return out


def row_lookup_to_json(row_lookup: dict[tuple[str, int], tuple[str, str]]) -> str:
    """Serialize the (session_date, row_index) -> (turn_doc_id, problem_doc_id) map.

    JSON doesn't support tuple keys, so we flatten to a list of records.
    """
    flat = [
        {
            "session_date": k[0],
            "row_index": k[1],
            "lecture_turn_doc_id": v[0],
            "lecture_problem_doc_id": v[1],
        }
        for k, v in sorted(row_lookup.items())
    ]
    return json.dumps(flat, indent=2)


def row_lookup_from_json(text: str) -> dict[tuple[str, int], tuple[str, str]]:
    """Deserialize the lookup map produced by row_lookup_to_json."""
    flat = json.loads(text)
    return {
        (rec["session_date"], rec["row_index"]):
            (rec["lecture_turn_doc_id"], rec["lecture_problem_doc_id"])
        for rec in flat
    }