"""Original LockedIn Labs tokenization exercise. Synthetic text; no model calls.

Python 3.11+. Install the companion requirements, then run `prepare` once.
Preparation downloads the two public encoding dictionaries. Other commands
block socket connections and use the local cache. Counts cover plain text only.
"""
from __future__ import annotations

import argparse
import importlib.metadata
import json
import os
from pathlib import Path
import socket
import unittest
from unittest.mock import patch

import tiktoken

ENCODINGS = ("cl100k_base", "o200k_base")
SAMPLES = (
    ("sentence", "A service request", "The pump needs inspection."),
    ("code", "A little code", 'ticket_id = 104\nstatus = "open"'),
    ("japanese", "The same kind of request in Japanese", "ポンプを点検してください。"),
)
os.environ.setdefault("TIKTOKEN_CACHE_DIR", str(Path.cwd() / ".tokenizer-cache"))

class EncodingCacheNeeded(RuntimeError):
    """The learner has not completed the explicit network setup step."""


def inspect_text(text: str, encoding_name: str) -> dict:
    encoding = tiktoken.get_encoding(encoding_name)
    # Treat all supplied text literally, including special-looking strings.
    ids = encoding.encode_ordinary(text)
    pieces = []
    for token_id in ids:
        raw = encoding.decode_single_token_bytes(token_id)
        try:
            readable = raw.decode("utf-8")
        except UnicodeDecodeError:
            readable = None
        pieces.append({"id": token_id, "bytes": list(raw), "text": readable})
    return {
        "encoding": encoding_name,
        "text": text,
        "codepoints": len(text),
        "utf8_bytes": len(text.encode("utf-8")),
        "tokens": pieces,
        "roundtrip": encoding.decode(ids) == text,
    }


def budget(text_tokens: int, history: int = 0, output: int = 48) -> dict:
    """A deliberately tiny, hypothetical window, not a model specification."""
    if any(type(value) is not int or value < 0 for value in (text_tokens, history, output)):
        raise ValueError("Budgets must be non-negative whole numbers.")
    window, scaffold = 128, 24
    input_count = text_tokens + history + scaffold
    total = input_count + output
    return {
        "window": window, "scaffold": scaffold, "text": text_tokens,
        "history": history, "output": output, "input": input_count,
        "reserved": total, "remaining": max(0, window - total),
        "over_by": max(0, total - window), "fits": total <= window,
    }


def recordings() -> dict:
    return {
        "runtime": {"tiktoken": importlib.metadata.version("tiktoken")},
        "samples": [
            {"id": key, "label": label, "encodings": [inspect_text(text, name) for name in ENCODINGS]}
            for key, label, text in SAMPLES
        ],
    }


class TokenTests(unittest.TestCase):
    def test_every_sample_roundtrips_from_exact_bytes(self):
        for _, _, text in SAMPLES:
            for name in ENCODINGS:
                run = inspect_text(text, name)
                self.assertTrue(run["roundtrip"])
                self.assertEqual(b"".join(bytes(t["bytes"]) for t in run["tokens"]), text.encode())

    def test_tokens_are_not_codepoints_or_bytes(self):
        run = inspect_text(SAMPLES[0][2], ENCODINGS[0])
        self.assertNotEqual(len(run["tokens"]), run["codepoints"])
        self.assertNotEqual(len(run["tokens"]), run["utf8_bytes"])

    def test_encoding_changes_the_japanese_representation(self):
        left, right = (inspect_text(SAMPLES[2][2], name) for name in ENCODINGS)
        self.assertNotEqual(left["tokens"], right["tokens"])
        self.assertNotEqual(len(left["tokens"]), len(right["tokens"]))

    def test_a_token_can_be_only_part_of_a_utf8_character(self):
        run = inspect_text(SAMPLES[2][2], ENCODINGS[0])
        self.assertTrue(any(token["text"] is None for token in run["tokens"]))

    def test_literal_special_looking_text_is_preserved(self):
        self.assertTrue(inspect_text("<|endoftext|>", ENCODINGS[0])["roundtrip"])
        self.assertEqual(inspect_text("", ENCODINGS[0])["tokens"], [])

    def test_exact_fit_and_one_over_are_distinct(self):
        self.assertEqual(budget(8, 80, 16)["reserved"], 128)
        self.assertTrue(budget(8, 80, 16)["fits"])
        self.assertFalse(budget(9, 80, 16)["fits"])
        self.assertEqual(budget(9, 80, 16)["over_by"], 1)

    def test_output_is_reserved_not_measured(self):
        run = budget(5, 80, 48)
        self.assertEqual(run["input"], 109)
        self.assertEqual(run["reserved"], 157)
        self.assertEqual(run["over_by"], 29)
        self.assertNotIn("usage", run)

    def test_invalid_budget_is_rejected(self):
        for value in (-1, 1.5, True):
            with self.assertRaises(ValueError):
                budget(value)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("prepare", "demo", "inspect", "budget", "test"))
    parser.add_argument("--text", default=SAMPLES[0][2], help="Use synthetic text only.")
    parser.add_argument("--encoding", choices=ENCODINGS, default=ENCODINGS[0])
    parser.add_argument("--history", type=int, default=0)
    parser.add_argument("--output", type=int, default=48)
    args = parser.parse_args()
    if args.command == "prepare":
        for name in ENCODINGS:
            tiktoken.get_encoding(name)
        print("Both public encoding dictionaries are cached. Subsequent commands block network access.")
        return
    # Missing cache is an explicit setup error, never a hidden network fetch.
    offline = EncodingCacheNeeded("Network blocked. Run prepare first in the same folder.")
    with patch.object(socket.socket, "connect", side_effect=offline), patch.object(socket.socket, "connect_ex", side_effect=offline), patch.object(socket, "getaddrinfo", side_effect=offline):
        if args.command == "test":
            result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(TokenTests))
            raise SystemExit(not result.wasSuccessful())
        if args.command == "demo":
            result = recordings()
        else:
            result = inspect_text(args.text, args.encoding)
            if args.command == "budget":
                try:
                    result = {"encoding": args.encoding, "text": args.text, "budget": budget(len(result["tokens"]), args.history, args.output)}
                except ValueError as error:
                    parser.error(str(error))
        print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    try:
        main()
    except EncodingCacheNeeded as error:
        raise SystemExit(str(error)) from None