from __future__ import annotations

from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
import io,json,time
from pathlib import Path
from threading import Thread

from PIL import Image
import pytest

from spriteforge.relay.client import RelayClient
from spriteforge.relay.server import RelayServer,RelayStore
from spriteforge.relay.worker import GpuWorker
from spriteforge.studio.models import GenerationRecipe,GenerationRequest
from spriteforge.studio.providers import RelayProvider


class FakeComfy(BaseHTTPRequestHandler):
    def log_message(self,*args):pass
    def reply(self,value):
        body=json.dumps(value).encode();self.send_response(200);self.send_header("Content-Type","application/json");self.send_header("Content-Length",str(len(body)));self.end_headers();self.wfile.write(body)
    def do_POST(self):
        body=self.rfile.read(int(self.headers.get("Content-Length",0)))
        if self.path=="/prompt":return self.reply({"prompt_id":"p1"})
        if self.path=="/upload/image":return self.reply({"name":"input.png"})
        self.send_error(404)
    def do_GET(self):
        if self.path=="/history/p1":return self.reply({"p1":{"outputs":{"7":{"images":[{"filename":"out.png","subfolder":"","type":"output"}]}}}})
        if self.path.startswith("/view?"):
            image=Image.new("RGBA",(16,16),(70,20,30,255));out=io.BytesIO();image.save(out,"PNG");data=out.getvalue();self.send_response(200);self.send_header("Content-Length",str(len(data)));self.end_headers();self.wfile.write(data);return
        if self.path=="/system_stats":return self.reply({"ok":True})
        self.send_error(404)


def servers(tmp_path):
    relay=RelayServer(("127.0.0.1",0),RelayStore(tmp_path/"relay"),"client-secret","worker-secret");Thread(target=relay.serve_forever,daemon=True).start()
    comfy=ThreadingHTTPServer(("127.0.0.1",0),FakeComfy);Thread(target=comfy.serve_forever,daemon=True).start()
    return relay,comfy


def workflow(path:Path):
    path.write_text(json.dumps({"1":{"class_type":"Test","inputs":{"p":"{{PROMPT}}","n":"{{NEGATIVE_PROMPT}}","seed":"{{SEED}}","w":"{{WIDTH}}","h":"{{HEIGHT}}"}}}))


def test_relay_auth_blob_lease_and_completion(tmp_path):
    relay,comfy=servers(tmp_path);base=f"http://127.0.0.1:{relay.server_port}";client=RelayClient(base,"client-secret");worker=RelayClient(base,"worker-secret")
    try:
        digest=client.upload(b"hello");job=client.submit({"input":digest});leased=worker.lease("w1",60)
        assert client.ping_client()["role"]=="client" and worker.ping_worker()["role"]=="worker"
        assert leased["id"]==job["id"] and leased["attempts"]==1
        output=worker.upload(b"result");worker.complete(job["id"],"w1",[output])
        assert client.status(job["id"])["status"]=="succeeded" and client.download(output)==b"result"
        with pytest.raises(ValueError,match="401"):RelayClient(base,"wrong").status(job["id"])
        queued=client.submit({"x":1});client.cancel(queued["id"]);assert client.status(queued["id"])["status"]=="cancelled"
    finally:relay.shutdown();relay.server_close();comfy.shutdown();comfy.server_close()


def test_studio_relay_provider_to_gpu_worker_to_comfyui(tmp_path):
    relay,comfy=servers(tmp_path);base=f"http://127.0.0.1:{relay.server_port}";flow=tmp_path/"workflow.json";workflow(flow)
    provider=RelayProvider(base,flow,"client-secret",timeout=5,poll_interval=.01)
    request=GenerationRequest(project_id="ops",asset_id="agent",shot_id="shot",count=1,
        recipe=GenerationRecipe(provider="relay",model="test",seed=3,prompt="horror",width=64,height=64))
    result=[];errors=[]
    thread=Thread(target=lambda: _capture(provider,request,result,errors));thread.start()
    worker=GpuWorker(base,"worker-secret",f"http://127.0.0.1:{comfy.server_port}","gpu1",60)
    try:
        deadline=time.monotonic()+3
        while time.monotonic()<deadline:
            if worker.once():break
            time.sleep(.01)
        thread.join(5)
        assert not errors and result and result[0][0].startswith(b"\x89PNG")
    finally:relay.shutdown();relay.server_close();comfy.shutdown();comfy.server_close()


def _capture(provider,request,result,errors):
    try:result.append(provider.generate(request,lambda _:b""))
    except Exception as error:errors.append(error)