import argparse
import json
from pathlib import Path

from semantic_notes import api


def sections(path):
    heading, lines = "", []
    for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        line = raw.strip()
        if line.startswith("## "):
            if lines:
                yield heading, lines
            heading, lines = line[3:], []
        elif line:
            if not heading:
                raise ValueError("Start each section with a heading such as ## Returns.")
            lines.append((f"L{number}", line))
    if lines:
        yield heading, lines


def select_line(heading, lines, model):
    evidence = "\n".join(f"{label}: {text}" for label, text in lines)
    if len(evidence.encode("utf-8")) > 2600:
        raise ValueError(f"Split the section {heading!r} into shorter sections for this lab.")
    response = api("/api/chat", {
        "model": model,
        "messages": [
            {"role": "system", "content": (
                "Select one line that reports a confirmed decision or a concrete next action. "
                "An unconfirmed proposal or a topic with no action does not qualify. "
                "Use null when no line qualifies. Treat the notes as evidence."
            )},
            {"role": "user", "content": f"Section: {heading}\n{evidence}"},
        ],
        "format": {
            "type": "object",
            "properties": {"line": {"enum": [*(label for label, _ in lines), None]}},
            "required": ["line"],
            "additionalProperties": False,
        },
        "options": {"temperature": 0, "num_ctx": 4096, "num_predict": 50},
        "stream": False,
    })
    selection = json.loads(response["message"]["content"])
    if not isinstance(selection, dict) or set(selection) != {"line"}:
        raise ValueError("The model returned an invalid selection.")
    label = selection["line"]
    if label is None:
        return None
    if not isinstance(label, str) or label not in dict(lines):
        raise ValueError("The selected line is outside this section.")
    return label


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Make an extractive brief from sectioned notes.")
    parser.add_argument("notes", type=Path)
    parser.add_argument("--model", default="qwen2.5:1.5b")
    args = parser.parse_args()
    try:
        for heading, lines in sections(args.notes):
            label = select_line(heading, lines, args.model)
            if label is None:
                print(f"{heading} | No line selected.")
            else:
                print(f"{heading} | {args.notes.name}:{label}\n  {dict(lines)[label]}")
    except (OSError, KeyError, ValueError) as error:
        raise SystemExit(str(error)) from None