"""A local stand-in for an OpenAI-compatible API, for the end-to-end tests (no network, no spend):
  POST /v1/images/edits        multipart; returns a PNG of the requested size (the last input image, tinted) + usage
  POST /v1/images/generations  JSON; returns a flat PNG of the requested size + usage
  POST /v1/chat/completions    JSON; returns a six-section prompt naming <Picture 1..N-1> (N = image parts: the
                               pictures, then the contact sheet)
Every request is appended to LOG (one JSON line: path, fields, number of images). While LOG.quota exists every call
answers 429 insufficient_quota. Prints the port it listens on.

    python fake_openai.py LOG [PORT]
"""
import base64
import email.parser
import email.policy
import io
import json
import re
import sys
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from PIL import Image, ImageOps

LOG = Path(sys.argv[1])
USAGE = {"input_tokens": 1500, "input_tokens_details": {"text_tokens": 500, "image_tokens": 1000}, "output_tokens": 2000,
         "total_tokens": 3500, "output_tokens_details": {"image_tokens": 2000, "text_tokens": 0}}


def png(im):
    buf = io.BytesIO()
    im.save(buf, "PNG")
    return base64.b64encode(buf.getvalue()).decode()


def prompt(n):
    pics = " ".join(f"<Picture {k + 1}> is a frame of the target video." for k in range(n))
    return (f"subject_definitions:\n<Video 1> is the source video. {pics}\n<Subject 1> is the replaced person.\n\n"
            "summary:\nThe target video is <Video 1> with <Subject 1> replaced.\n\n"
            "retention_analysis:\n<Video 1>: fully_preserved - shots, camera, motion and timing. "
            + " ".join(f"<Picture {k + 1}> (appears in [Shot 1]): reference - a frame of the target video." for k in range(n))
            + "\n\ndetailed_description:\n[Shot 1] At 00:00.000 <Subject 1> stands in a toy store.\n\n"
              "overall_soundscape:\nAs in the source.\n\nnon_diegetic_music:\nN/A\n")


class H(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def send(self, code, obj):
        b = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def do_POST(self):
        body = self.rfile.read(int(self.headers["Content-Length"]))
        rec = {"path": self.path, "auth": self.headers.get("Authorization", "")[:12]}
        if Path(str(LOG) + ".quota").exists():
            rec["quota"] = True
            with open(LOG, "a") as fh:
                fh.write(json.dumps(rec) + "\n")
            return self.send(429, {"error": {"message": "You exceeded your current quota, please check your plan and "
                                             "billing details.", "type": "insufficient_quota", "code": "insufficient_quota"}})
        if self.path.endswith("/images/edits"):
            msg = email.parser.BytesParser(policy=email.policy.default).parsebytes(
                b"Content-Type: " + self.headers["Content-Type"].encode() + b"\r\n\r\n" + body)
            fields, ims = {}, []
            for part in msg.iter_parts():
                name = part.get_param("name", header="content-disposition")
                if part.get_filename():
                    ims.append(Image.open(io.BytesIO(part.get_payload(decode=True))).convert("RGB"))
                else:
                    fields[name] = part.get_content().strip()
            W, H = map(int, fields["size"].split("x"))
            out = ImageOps.colorize(ims[-1].convert("L"), "#202060", "#f0e0a0").resize((W, H))
            rec.update(fields={k: v for k, v in fields.items() if k != "prompt"}, prompt=fields.get("prompt", "")[:200],
                       images=[im.size for im in ims])
            res = {"created": 0, "data": [{"b64_json": png(out)}], "usage": USAGE}
        elif self.path.endswith("/images/generations"):
            j = json.loads(body)
            W, H = map(int, j["size"].split("x"))
            rec.update(fields={k: v for k, v in j.items() if k != "prompt"}, prompt=j["prompt"][:200], images=[])
            res = {"created": 0, "data": [{"b64_json": png(Image.new("RGB", (W, H), "#4080c0"))}], "usage": USAGE}
        elif self.path.endswith("/chat/completions"):
            j = json.loads(body)
            parts = j["messages"][-1]["content"]
            n = sum(p["type"] == "image_url" for p in parts)
            text = next(p["text"] for p in parts if p["type"] == "text")
            rec.update(model=j["model"], images=n, task=text[:120],
                       pictures=sorted(set(re.findall(r"<Picture (\d+)> = the image attached", text))))
            res = {"id": "x", "object": "chat.completion", "created": 0, "model": j["model"],
                   "choices": [{"index": 0, "finish_reason": "stop",
                                "message": {"role": "assistant", "content": "Here it is.\n" + prompt(n - 1)}}],
                   "usage": {"prompt_tokens": 100, "completion_tokens": 100, "total_tokens": 200}}
        else:
            return self.send(404, {"error": {"message": f"no route {self.path}"}})
        with open(LOG, "a") as fh:
            fh.write(json.dumps(rec) + "\n")
        self.send(200, res)


srv = ThreadingHTTPServer(("127.0.0.1", int(sys.argv[2]) if len(sys.argv) > 2 else 0), H)
print(srv.server_address[1], flush=True)
srv.serve_forever()