"""Pack related decisions into bounded multi-question API requests."""

import json
import sqlite3

from tokenizers import Tokenizer

from .schema import digest


def pack_release(root, version, cfg, lock):
    out = root / "data/curated" / version
    tm = lock["tokenizer"]
    tokenizer = Tokenizer.from_file(
        str(root / "data/raw" / tm["repo"].replace("/", "--") / "tokenizer.json")
    )
    tokenizer.no_truncation()
    tokenizer.no_padding()
    budget = cfg["max_input_tokens"] - cfg["prompt_reserve_tokens"]
    counts = {}
    for split in ("train", "validation", "calibration", "test"):
        # Disk-backed grouping avoids retaining whole source texts in memory.
        db_path = root / "data/work" / ("pack-" + split + ".sqlite")
        if db_path.exists():
            db_path.unlink()
        db = sqlite3.connect(db_path)
        db.execute("CREATE TABLE records (state_key TEXT, id TEXT, body TEXT)")
        with (out / (split + ".jsonl")).open() as stream:
            for line in stream:
                record = json.loads(line)
                db.execute(
                    "INSERT INTO records VALUES (?,?,?)",
                    (digest(record["input"]["state"]), record["id"], line),
                )
        count, decisions, multi = 0, 0, 0
        current, state_key = None, None
        with (out / (split + "-requests.jsonl")).open("w") as stream:

            def write(request):
                nonlocal count, decisions, multi
                if request is None or not request["targets"]:
                    return
                n = len(request["targets"])
                count += 1
                decisions += n
                multi += int(n > 1)
                stream.write(json.dumps(request, ensure_ascii=False, separators=(",", ":")) + "\n")

            for key, _, body in db.execute("SELECT * FROM records ORDER BY state_key,id"):
                record = json.loads(body)
                if state_key != key:
                    write(current)
                    current, state_key = None, key
                if current is None:
                    current = {
                        "input": {"state": record["input"]["state"], "questions": {}},
                        "targets": {},
                        "option_keys": {},
                        "decision_ids": {},
                        "metadata": {
                            "split": split,
                            "state_hash": record["metadata"]["state_hash"],
                        },
                    }
                qid = "decision" if not current["targets"] else "q" + str(len(current["targets"]))
                current["input"]["questions"][qid] = record["input"]["questions"]["decision"]
                serialized = json.dumps(current["input"], ensure_ascii=False, separators=(",", ":"))
                n_tokens = len(tokenizer.encode(serialized, add_special_tokens=False).ids)
                if n_tokens > budget:
                    del current["input"]["questions"][qid]
                    write(current)
                    current = {
                        "input": {
                            "state": record["input"]["state"],
                            "questions": {"decision": record["input"]["questions"]["decision"]},
                        },
                        "targets": {},
                        "option_keys": {},
                        "decision_ids": {},
                        "metadata": {
                            "split": split,
                            "state_hash": record["metadata"]["state_hash"],
                        },
                    }
                    qid = "decision"
                    serialized = json.dumps(
                        current["input"], ensure_ascii=False, separators=(",", ":")
                    )
                    n_tokens = len(tokenizer.encode(serialized, add_special_tokens=False).ids)
                    if n_tokens > budget:
                        raise ValueError("Single decision no longer fits request budget")
                current["targets"][qid] = record["target"]
                current["option_keys"][qid] = record["option_keys"]
                current["decision_ids"][qid] = record["id"]
                current["metadata"]["payload_tokens"] = n_tokens
            write(current)
        counts[split] = {
            "requests": count,
            "decisions": decisions,
            "multi_question_requests": multi,
        }
        db.close()
        db_path.unlink()
    return counts