import argparse
import json
from pathlib import Path

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

from semantic_notes import api

LABELS = ("RETURN", "RENEW", "RESERVE")
REVIEW = "REVIEW"
OTHER_REQUESTS = {
    "valid": [
        ("Does the library sell headphones?", REVIEW),
        ("Can I return this book and book the projector?", REVIEW),
        ("Where are the study rooms?", REVIEW),
    ],
    "test": [
        ("What time does the library open?", REVIEW),
        ("I would like to buy a notebook.", REVIEW),
        ("Could I book the scanner and extend my book loan?", REVIEW),
    ],
}


def load_examples(path):
    examples = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            system, user, assistant = json.loads(line)["messages"]
            examples.append((system["content"], user["content"], assistant["content"]))
    if not examples:
        raise ValueError(f"No examples in {path}.")
    return examples


def ask_chat(instructions, text, allowed, constrained):
    payload = {
        "model": "qwen2.5:1.5b",
        "messages": [
            {"role": "system", "content": instructions},
            {"role": "user", "content": text},
        ],
        "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 30},
        "stream": False,
    }
    if constrained:
        payload["format"] = {
            "type": "object",
            "properties": {"label": {"enum": allowed}},
            "required": ["label"],
            "additionalProperties": False,
        }
    content = api("/api/chat", payload)["message"]["content"].strip()
    if constrained:
        try:
            content = json.loads(content)["label"]
        except (KeyError, TypeError, ValueError):
            return "(invalid)"
    return content if isinstance(content, str) else "(invalid)"


def report(title, predictions, expected):
    correct = sum(guess == target for guess, target in zip(predictions, expected))
    reviewed = sum(guess == REVIEW for guess in predictions)
    print(f"{title}: {correct}/{len(expected)} correct, {reviewed} sent to review")
    for label in (*LABELS, REVIEW):
        count = sum(target == label for target in expected)
        if count:
            matches = sum(guess == target == label for guess, target in zip(predictions, expected))
            print(f"  {label}: {matches}/{count}")


def compare(args):
    train = load_examples(args.data / "train.jsonl")
    data = load_examples(args.data / f"{args.split}.jsonl")
    instructions = data[0][0]
    samples = [(user, label) for _, user, label in data]
    allowed = list(LABELS)
    if args.include_review:
        samples.extend(OTHER_REQUESTS[args.split])
        allowed.append(REVIEW)
        instructions = (
            instructions.replace("Reply with exactly one label.", "") +
            " Use REVIEW for other requests or for requests with multiple intents. "
            "Reply with exactly one label."
        )
    texts, expected = zip(*samples)
    pipeline = make_pipeline(
        TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True),
        LogisticRegression(C=4, max_iter=1000),
    )
    pipeline.fit([user for _, user, _ in train], [label for _, _, label in train])
    probabilities = pipeline.predict_proba(texts)
    predictions = []
    for row, target, text in zip(probabilities, expected, texts):
        index = row.argmax()
        score = row[index]
        label = str(pipeline.classes_[index]) if score >= args.threshold else REVIEW
        predictions.append(label)
        print(f"{target:7} | {label:7} {score:.2f} | {text}")
    report("Classifier", predictions, expected)
    if args.compare_chat:
        guesses = [ask_chat(instructions, text, allowed, args.json_chat) for text in texts]
        for target, guess, text in zip(expected, guesses, texts):
            print(f"{target:7} | {guess!r} | {text}")
        report("Chat model", guesses, expected)
        print("  Invalid labels:", sum(guess not in allowed for guess in guesses))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Compare a small classifier and an optional local chat baseline.")
    parser.add_argument("--data", type=Path, default=Path("library_data"))
    parser.add_argument("--split", choices=("valid", "test"), default="test")
    parser.add_argument("--threshold", type=float, default=0.0)
    parser.add_argument("--include-review", action="store_true")
    parser.add_argument("--compare-chat", action="store_true")
    parser.add_argument("--json-chat", action="store_true")
    args = parser.parse_args()
    if not 0 <= args.threshold <= 1:
        parser.error("Choose a threshold from 0 to 1.")
    if args.json_chat and not args.compare_chat:
        parser.error("--json-chat requires --compare-chat.")
    try:
        compare(args)
    except (OSError, KeyError, TypeError, ValueError) as error:
        raise SystemExit(str(error)) from None