"""Task 1: classify which oumi subsystem a PR touches, from its title + body."""

import json
import random
import re

from app.config import DATA_DIR
from app.tasks.base import Sample, Task

AREAS = [
    "training", "inference", "data", "evaluation", "configs",
    "launcher", "cli", "docs", "infra", "other",
]

SYSTEM_PROMPT = """\
You are a pull-request triage classifier for oumi, an open-source library for \
training, evaluating, and deploying foundation models.

Given a PR's title and body, classify which subsystem of the repo it primarily \
changes. The labels:

- training: training loop, trainers, tuners, callbacks, collators, distributed/performance code
- inference: inference engines, serving, deployment code
- data: dataset classes, tokenizers, processors, data synthesis, conversation types
- evaluation: evaluation harnesses, judges, analyzers, metrics
- configs: YAML recipes for models/jobs, or the config schema code (params classes)
- launcher: launching/managing jobs on clouds and clusters (SkyPilot, Slurm, ...)
- cli: the `oumi` command-line interface
- docs: documentation, notebooks, markdown files
- infra: CI/CD, GitHub workflows, dependency bumps, build tooling, repo maintenance
- other: cross-cutting utilities, model architecture code, or anything not above

Respond with exactly one JSON object on a single line, no other text:
{"area": "<label>"}\
"""

_RECORDS: dict[str, dict] = {}
with (DATA_DIR / "prs.jsonl").open() as _f:
    for _line in _f:
        _r = json.loads(_line)
        _RECORDS[f"pr-{_r['number']}"] = _r

_rng = random.Random()


def render_pr(record: dict) -> str:
    body = (record.get("body") or "").strip()
    text = f"Title: {record['title']}"
    if body:
        text += f"\n\nBody:\n{body}"
    return text


def sample() -> Sample:
    input_id = _rng.choice(list(_RECORDS))
    return Sample(input_id=input_id, text=render_pr(_RECORDS[input_id]))


def lookup_truth(input_id: str) -> dict | None:
    record = _RECORDS.get(input_id)
    if record is None:
        return None
    return {"area": record["area"], "pr_number": record["number"]}


def parse_output(text: str) -> dict:
    m = re.search(r"\{[^{}]*\}", text, re.DOTALL)
    area = None
    if m:
        try:
            value = json.loads(m.group()).get("area")
            if isinstance(value, str) and value.strip().lower() in AREAS:
                area = value.strip().lower()
        except json.JSONDecodeError:
            pass
    return {"area": area, "raw": text}


def score(truth: dict, parsed: dict) -> dict[str, float]:
    return {
        "format_ok": float(parsed["area"] is not None),
        "exact_match": float(parsed["area"] == truth["area"]),
    }


def present(parsed: dict, truth: dict | None) -> dict:
    return {"area": parsed["area"], "raw": parsed["raw"]}


TASK = Task(
    order=1,
    id="pr-area",
    title="PR Triage",
    tagline="Which oumi subsystem does this PR touch?",
    system_prompt=SYSTEM_PROMPT,
    ui={
        "output": "label",
        "labels": AREAS,
        "input_label": "Pull request (title + body)",
        "placeholder": "Deal a real oumi PR, or write your own title + body...",
        "deal_label": "Deal me a PR",
    },
    sample=sample,
    lookup_truth=lookup_truth,
    parse_output=parse_output,
    score=score,
    present=present,
)