"""Run with pinned gzipt checkout as the first argument. No third-party deps."""
import collections
import importlib.util
import json
from pathlib import Path
import random
import subprocess
import sys
import tempfile
import zlib

source = Path(sys.argv[1]).resolve() / "gzipt.py"
spec = importlib.util.spec_from_file_location("gzipt", source)
gzipt = importlib.util.module_from_spec(spec)
spec.loader.exec_module(gzipt)
print(json.dumps({"python": sys.version.split()[0], "zlib": zlib.ZLIB_RUNTIME_VERSION}))
rng = random.Random(17)
contexts = [b"", b"abc " * 100, rng.randbytes(40000)]
sequences = [b"", b"a", b"abc " * 12, rng.randbytes(97)]
checks = 0
mismatches = []
for level in (0, 1, 6, 9):
    for context in contexts:
        actual = gzipt.candidate_lengths(context, sequences, level=level)
        expected = [len(zlib.compress(context + seq, level)) for seq in sequences]
        for i, (a, e) in enumerate(zip(actual, expected)):
            if a != e:
                mismatches.append({"level": level, "context_bytes": len(context), "sequence_bytes": len(sequences[i]), "clone": a, "full": e})
        checks += len(sequences)
print(json.dumps({"compression_comparisons": checks, "mismatches": mismatches}))
corpus = b"the cat sat on the mat. the dog sat on the rug.\n" * 80
prompt = b"the "
alphabet = gzipt.corpus_alphabet(corpus + prompt)
ctx = corpus + prompt
scores = gzipt.candidate_lengths(ctx, [bytes([b]) for b in alphabet])
best = min(scores)
print(json.dumps({"alphabet_size": len(alphabet), "score_histogram": dict(sorted(collections.Counter(scores).items())), "best_bytes": bytes(b for b, score in zip(alphabet, scores) if score == best).decode()}))
for horizon in (1, 4, 8):
    result = gzipt.generate(corpus, prompt, 32, horizon=horizon, beam_width=16, temperature=0, workers=1)
    print(json.dumps({"horizon": horizon, "output": result.decode()}))

# Observe the contexts passed by the actual generator, without changing scoring.
original = gzipt.candidate_lengths
for tail in (0, 1, 4):
    observed = []
    def capture(context, sequences, **kwargs):
        observed.append(context)
        return original(context, sequences, **kwargs)
    gzipt.candidate_lengths = capture
    gzipt.generate(b"abc", b"PROMPT", 2, window=3, horizon=1, beam_width=2, temperature=0, tail=tail, alphabet=(97,98))
    print(json.dumps({"tail": tail, "first_context": observed[0].decode(), "second_context": observed[1].decode()}))
gzipt.candidate_lengths = original

with tempfile.TemporaryDirectory() as folder:
    fixture = Path(folder) / "corpus.txt"
    fixture.write_bytes(corpus)
    command = [sys.executable, str(source), "--corpus", str(fixture), "--prompt", "the ", "--length", "32", "--horizon", "4", "--beam-width", "16", "--temperature", "0", "--workers", "1"]
    run = subprocess.run(command, capture_output=True, check=True)
    print(json.dumps({"cli_exit": run.returncode, "stdout": run.stdout.decode().replace("\r\n", "\n")}))