#!/usr/bin/env python3
"""Work on zetatron problems from the terminal."""

import argparse
import html
import json
import os
import re
import string
import sys
import urllib.error
import urllib.request

# Replaced with the server's own address when downloaded from it.
SERVER = "https://zetatron.thinkinrocks.com"
PINCODE_FILE = ".pincode"
PINCODE_LENGTH = 16
# Set to "true" to agree that no generative AI was used without being asked.
NO_GEN_AI = "ZETATRON_NO_GEN_AI_USED"


def fetch(path):
    """The body of `path` on the server, or `None` if it isn't there."""
    try:
        with urllib.request.urlopen(SERVER + path) as response:
            return response.read()
    except urllib.error.HTTPError as error:
        if error.code == 404:
            return None
        raise


def text(description):
    """A test's HTML description as plain text."""
    return html.unescape(re.sub(r"<[^>]+>", "", description)).strip()


def setup(_):
    me = os.path.basename(sys.argv[0])
    if [name for name in os.listdir() if name != me]:
        answer = input("This folder is not empty, and setup will add files to it. Continue? Type yes: ")
        if answer.strip().lower() != "yes":
            sys.exit("Stopped.")

    paths = ["/.clang-format"]
    for problem in json.loads(fetch("/p.json")):
        paths += [f"/p/{problem['name']}.md", f"/p/{problem['name']}.cpp"]
    for path in paths:
        name = os.path.basename(path)
        if os.path.exists(name):
            print(f"Skipped {name}, it already exists.")
            continue
        body = fetch(path)
        if body is not None:
            with open(name, "wb") as file:
                file.write(body)
            print(f"Downloaded {name}.")

    while True:
        pincode = input(f"Your {PINCODE_LENGTH}-character pincode: ").strip().upper()
        if len(pincode) == PINCODE_LENGTH and all(c in string.ascii_uppercase + string.digits for c in pincode):
            break
        print(f"A pincode is {PINCODE_LENGTH} letters and digits.")
    with open(PINCODE_FILE, "w") as file:
        file.write(pincode + "\n")
    print(f"Saved the pincode to {PINCODE_FILE}.")


def submit(args):
    names = [problem["name"] for problem in json.loads(fetch("/p.json"))]
    matches = [name for name in names if name == args.problem] or [
        name for name in names if name.startswith(args.problem)
    ]
    if len(matches) != 1:
        sys.exit(f"{args.problem!r} matches {', '.join(matches) or 'no problem'}.")
    name = matches[0]
    try:
        with open(PINCODE_FILE) as file:
            pincode = file.read().strip()
        with open(f"{name}.cpp") as file:
            source = file.read()
    except FileNotFoundError as error:
        sys.exit(f"{error.filename} is missing, run setup first.")

    if os.environ.get(NO_GEN_AI) != "true":
        answer = input(f"Submitting {name}.cpp. Did you write it without any LLMs or other generative AI? Type yes: ")
        if answer.strip().lower() != "yes":
            sys.exit("Stopped.")
        print(f"To agree for the rest of this session, run: export {NO_GEN_AI}=true")

    request = urllib.request.Request(
        f"{SERVER}/p/{name}/submit",
        data=json.dumps({"source": source, "no_ai": True}).encode(),
        headers={"Content-Type": "application/json", "Cookie": f"zetatron_pincode={pincode}"},
    )
    try:
        with urllib.request.urlopen(request) as response:
            id = json.load(response)["id"]
    except urllib.error.HTTPError as error:
        sys.exit(f"Not submitted: {error.read().decode() or error.reason}")
    print(f"Submitted, see {SERVER}/s/{id}")

    started = None
    with urllib.request.urlopen(f"{SERVER}/s/{id}/events") as events:
        for line in events:
            if not line.startswith(b"data:"):
                continue
            event = json.loads(line[5:])
            kind = event["kind"]
            if kind == "status":
                print(f"{event['status'].capitalize()}...")
            elif kind == "started":
                started = event
            elif kind == "compile_error":
                print(event["log"])
            elif kind == "correctness":
                result = "ok" if event["ok"] else "wrong answer"
                print(f"  {event['label']}: {result}")
                if not event["ok"] and started["description"]:
                    print(text(started["description"]))
                started = None
            elif kind == "benchmark":
                if not event["ok"]:
                    result = "time limit exceeded"
                else:
                    result = f"{event['time_ms']:.3f} ms after {event['runs']} runs"
                print(f"\r\033[K  {event['label']}: {result}", end="\n" if event["done"] else "", flush=True)
                if event["done"]:
                    started = None
            elif kind == "output":
                for stream in ("stdout", "stderr"):
                    if event[stream]:
                        print(f"{stream}:\n{event[stream]}")
            elif kind == "verdict":
                # A test that never finished, such as one that ran out of time.
                if started and started["test"] == "correctness":
                    print(f"  {started['label']}: {event['verdict']}")
                    if started["description"]:
                        print(text(started["description"]))
                print(event["verdict"].capitalize())
                break


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(required=True, metavar="command")
    commands.add_parser("setup", help="download the problems and save your pincode").set_defaults(run=setup)
    submit_parser = commands.add_parser("submit", aliases=["s"], help="submit a solution and show its results")
    submit_parser.add_argument("problem", help="the problem's name, or the start of it")
    submit_parser.set_defaults(run=submit)
    args = parser.parse_args()
    args.run(args)


if __name__ == "__main__":
    main()