"""Local HTTP API and workers for the Relay Lab demo."""

import argparse
import json
import logging
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from queue_store import Conflict, QueueStore


SCENARIOS = {"success", "transient", "crash", "permanent"}
PUBLIC = Path(__file__).with_name("public")
ASSETS = {
    "/": ("index.html", "text/html; charset=utf-8"),
    "/app.js": ("app.js", "text/javascript; charset=utf-8"),
    "/style.css": ("style.css", "text/css; charset=utf-8"),
}


def worker(store, stop):
    while not stop.is_set():
        try:
            job = store.claim(time.time(), lease_seconds=5)
            if not job:
                stop.wait(0.2)
                continue
            if stop.wait(0.35):
                return

            scenario, attempt = job["scenario"], job["attempts"]
            # Leave the first crash attempt leased, as if its worker disappeared.
            if scenario == "crash" and attempt == 1:
                continue
            success = scenario != "permanent" and not (scenario == "transient" and attempt < 3)
            store.finish(job, success, time.time())
        except Exception:
            logging.exception("Worker operation failed")
            stop.wait(1)


def handler_for(store):
    class Handler(BaseHTTPRequestHandler):
        def send_data(self, status, body, content_type):
            self.send_response(status)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def send_json(self, status, payload):
            self.send_data(status, json.dumps(payload).encode(), "application/json")

        def do_GET(self):
            if self.path == "/api/state":
                self.send_json(200, store.snapshot())
            elif self.path in ASSETS:
                name, content_type = ASSETS[self.path]
                self.send_data(200, (PUBLIC / name).read_bytes(), content_type)
            else:
                self.send_json(404, {"error": "Route not found."})

        def do_POST(self):
            if self.path != "/api/jobs":
                self.send_json(404, {"error": "Route not found."})
                return
            if self.headers.get_content_type() != "application/json":
                self.send_json(415, {"error": "Use application/json."})
                return
            try:
                length = int(self.headers.get("Content-Length", 0))
                if not 0 < length <= 2048:
                    raise ValueError
                payload = json.loads(self.rfile.read(length))
                key, scenario = payload["key"], payload["scenario"]
                if not isinstance(key, str) or not 1 <= len(key) <= 80:
                    raise ValueError
                if not isinstance(scenario, str) or scenario not in SCENARIOS:
                    raise ValueError
            except (KeyError, TypeError, ValueError):
                self.send_json(400, {"error": "Provide a valid key and scenario."})
                return
            try:
                job_id, created = store.enqueue(key, scenario, time.time())
            except Conflict as error:
                self.send_json(409, {"error": str(error)})
                return
            self.send_json(201 if created else 200, {"id": job_id, "created": created})

        def log_message(self, format, *args):
            if self.path != "/api/state":
                super().log_message(format, *args)

    return Handler


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", type=int, default=4176)
    parser.add_argument("--db", type=Path, default=Path("relay.db"))
    args = parser.parse_args()
    store = QueueStore(args.db)
    stop = threading.Event()
    with ThreadingHTTPServer(("127.0.0.1", args.port), handler_for(store)) as server:
        workers = [threading.Thread(target=worker, args=(store, stop)) for _ in range(2)]
        for thread in workers:
            thread.start()
        print(f"Relay Lab is running at http://127.0.0.1:{args.port}", flush=True)
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            pass
        finally:
            stop.set()
            for thread in workers:
                thread.join()


if __name__ == "__main__":
    main()