#!/usr/bin/env python3
"""Minimal fake AWS Batch API (rest-json, POST /v1/<op>) to reproduce a Nextflow nf-amazon poll-forever bug.

Implements just enough of RegisterJobDefinition / DescribeJobDefinitions /
SubmitJob / DescribeJobs for a real nf-amazon@3.4.4 AwsBatchTaskHandler to
submit a job successfully, then simulates a job aging out: after
VANISH_AFTER_S seconds since submit, DescribeJobs returns an empty `jobs`
list for that job id forever (the exact `{"jobs": []}` shape real AWS
returns for an aged-out job id).
"""
import json
import os
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

VANISH_AFTER_S = float(os.environ.get("VANISH_AFTER_S", "15"))
PORT = int(os.environ.get("PORT", "8080"))

_jobs: dict[str, float] = {}  # jobId -> submitted_at
_names: dict[str, str] = {}  # jobId -> jobName


class Handler(BaseHTTPRequestHandler):
    def _read_json(self) -> dict:
        length = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(length) if length else b"{}"
        try:
            return json.loads(raw or b"{}")
        except json.JSONDecodeError:
            return {}

    def _respond(self, status: int, body: dict) -> None:
        payload = json.dumps(body).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def log_message(self, fmt, *args):  # noqa: A002 - keep stdlib signature
        print(f"[fake-batch] {self.address_string()} {fmt % args}", flush=True)

    def do_POST(self) -> None:  # noqa: N802 - stdlib method name
        body = self._read_json()
        path = self.path

        if path == "/v1/describejobdefinitions":
            # Always empty -> forces the handler to register a fresh definition.
            self._respond(200, {"jobDefinitions": []})
            return

        if path == "/v1/registerjobdefinition":
            name = body.get("jobDefinitionName", "fake-job-def")
            self._respond(200, {
                "jobDefinitionName": name,
                "jobDefinitionArn": f"arn:aws:batch:us-east-1:000000000000:job-definition/{name}:1",
                "revision": 1,
            })
            return

        if path == "/v1/submitjob":
            job_id = str(uuid.uuid4())
            _jobs[job_id] = time.monotonic()
            name = body.get("jobName", "fake-job")
            _names[job_id] = name
            print(f"[fake-batch] SubmitJob -> jobId={job_id} name={name}", flush=True)
            self._respond(200, {
                "jobArn": f"arn:aws:batch:us-east-1:000000000000:job/{job_id}",
                "jobName": name,
                "jobId": job_id,
            })
            return

        if path == "/v1/describejobs":
            ids = body.get("jobs", [])
            jobs = []
            for job_id in ids:
                submitted_at = _jobs.get(job_id)
                if submitted_at is None:
                    continue  # unknown id -> also omitted, same as "vanished"
                name = _names.get(job_id, "")
                upper = name.upper()
                age = time.monotonic() - submitted_at
                # Only jobs named *LOST* ever vanish; others (e.g. a batched
                # sibling) stay RUNNING forever, so a batched describe can
                # come back non-empty while still missing the lost job.
                if "LOST" in upper and age >= VANISH_AFTER_S:
                    print(f"[fake-batch] DescribeJobs jobId={job_id} name={name} age={age:.1f}s "
                          f"-> VANISHED (omitted from response)", flush=True)
                    continue  # <-- the trigger: job silently absent from the response
                print(f"[fake-batch] DescribeJobs jobId={job_id} name={name} age={age:.1f}s -> RUNNING", flush=True)
                jobs.append({
                    "jobArn": f"arn:aws:batch:us-east-1:000000000000:job/{job_id}",
                    "jobName": name or "fake-job",
                    "jobId": job_id,
                    "jobQueue": "fake-queue",
                    "status": "RUNNING",
                    "startedAt": int(submitted_at * 1000),
                    "jobDefinition": "fake-job-def",
                })
            print(f"[fake-batch] DescribeJobs requested={ids} -> "
                  f"returning {len(jobs)}/{len(ids)} job(s)", flush=True)
            self._respond(200, {"jobs": jobs})
            return

        if path == "/v1/canceljob" or path == "/v1/terminatejob":
            self._respond(200, {})
            return

        print(f"[fake-batch] UNHANDLED {path} body={body}", flush=True)
        self._respond(200, {})


if __name__ == "__main__":
    print(f"[fake-batch] listening on 0.0.0.0:{PORT}, VANISH_AFTER_S={VANISH_AFTER_S}", flush=True)
    ThreadingHTTPServer(("0.0.0.0", PORT), Handler).serve_forever()