"""Loopback callback receiver with live carrier previews. Run with ComfyUI's Python."""

from __future__ import annotations

import argparse
import asyncio
from collections import deque
from contextlib import suppress
from datetime import datetime, timezone
from hashlib import sha256
from io import BytesIO
import json
import logging
from pathlib import Path
import sys
from uuid import uuid4

from aiohttp import web
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from chick.codec import CarrierCodecError, decode_bytes


PAGE = """<!doctype html><html lang="zh-CN"><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>NotifyMe 本地测试结果</title>
<style>
body{font:16px system-ui;margin:32px auto;max-width:860px;padding:0 20px;background:#f5f6fa;color:#18202b}
article{background:white;padding:20px;margin:16px 0;border:1px solid #e0e5ed;border-radius:10px}
pre{white-space:pre-wrap;overflow-wrap:anywhere}time,.muted{color:#536275;font-size:14px}
code{background:#e9edf4;padding:4px}video,img{display:block;max-width:100%;max-height:480px;margin-top:14px;background:#111}
video{width:100%;min-height:180px}audio{width:100%;margin-top:14px}
button{font:inherit;background:#245dba;color:white;border:0;border-radius:6px;padding:9px 16px;cursor:pointer}
button:disabled{opacity:.6;cursor:wait}.error{color:#a21d2a}a{color:#245dba}
</style>
<h1>本地测试结果</h1><p>回调地址：<code id="url"></code></p>
<p class="muted">测试模式：回调固定返回 HTTP 200，NotifyMe 可继续执行后续节点。处理详情见接收记录和响应正文。</p>
<p>点击开启监听后，新结果会自动追加，正在播放的视频会保持播放。</p>
<button id="listen" type="button" aria-pressed="false">开启监听</button> <span id="connection" role="status">尚未开启监听</span>
<p id="status" class="muted"></p>
<h2>回调内容</h2><div id="messages"><p id="messages-empty" class="muted">等待 NotifyMe 消息…</p></div>
<h2>解码媒体</h2><p class="muted">载体从指定的本地输出目录读取；回调包含任务 id 和文本 msg，两条分支独立执行。</p>
<div id="media"><p id="media-empty" class="muted">等待载体 PNG…</p></div>
<script src="/viewer.js" defer></script></html>"""


def now():
    return datetime.now(timezone.utc).isoformat()


class ReceiverState:
    def __init__(self, log_path=None, carrier_dir=None, media_dir=None, password=""):
        self.log_path = Path(log_path) if log_path is not None else None
        self.carrier_dir = Path(carrier_dir).resolve() if carrier_dir is not None else None
        self.media_dir = Path(media_dir or ".local/decoded-media").resolve()
        self.password = password
        self.messages = deque(maxlen=100)
        self.received = 0
        self.media = {}
        self.files = {}
        self.clients = set()
        self.scan_lock = asyncio.Lock()
        self.publish_lock = asyncio.Lock()
        self.revision = 0
        self.session = uuid4().hex
        if self.log_path is not None and self.log_path.is_file():
            with self.log_path.open(encoding="utf-8") as log:
                for index, line in enumerate(log):
                    try:
                        record = json.loads(line)
                        if not isinstance(record.get("msg"), str):
                            continue
                        if "event_id" not in record:
                            # Earlier logs used id for the receiver's event identity.
                            record["event_id"] = record.pop("id", sha256(f"{index}:{line}".encode()).hexdigest())
                            record["id"] = ""
                        self.messages.append(record)
                        self.received += 1
                    except (ValueError, AttributeError):
                        continue

    def snapshot(self):
        return {"ok": True, "service": "chicktools-callback-receiver",
                "session": self.session, "revision": self.revision,
                "received": self.received, "messages": list(self.messages),
                "media": list(self.media.values())}

    def receive(self, message, task_id=""):
        record = {"event_id": uuid4().hex, "id": task_id,
                  "received_at": now(), "path": "/callback", "msg": message}
        if self.log_path is not None:
            self.log_path.parent.mkdir(parents=True, exist_ok=True)
            with self.log_path.open("a", encoding="utf-8", newline="\n") as log:
                log.write(json.dumps(record, ensure_ascii=False) + "\n")
        self.messages.append(record)
        self.received += 1
        self.revision += 1

    def _scan(self):
        """Decode off the event loop; publish only complete, stable file versions."""
        records, files = [], {}
        if self.carrier_dir is None:
            return records, files
        self.media_dir.mkdir(parents=True, exist_ok=True)
        candidates = []
        for path in self.carrier_dir.glob("*.png"):
            try:
                if path.resolve().parent == self.carrier_dir:
                    stat = path.stat()
                    candidates.append((stat.st_mtime_ns, path, stat.st_size))
            except OSError:
                continue
        for modified, path, size in sorted(candidates)[-100:]:
            key = sha256(f"{path.name}:{modified}:{size}".encode()).hexdigest()
            if key in self.media:
                continue
            record = {"id": key, "source": path.name, "extension": "", "size": 0, "decoded_at": now()}
            try:
                source = path.read_bytes()
                current = path.stat()
                if (current.st_mtime_ns, current.st_size) != (modified, size):
                    continue
                with Image.open(BytesIO(source)) as image:
                    payload, extension = decode_bytes(image, self.password)
                extension = extension.lower().lstrip(".")
                if extension not in {"png", "jpg", "jpeg", "gif", "webp", "mp4", "webm", "wav", "mp3", "ogg"}:
                    raise CarrierCodecError("unsupported preview format")
                name = f"{key}.{extension}"
                target = self.media_dir / name
                temporary = self.media_dir / f"{key}.tmp"
                temporary.write_bytes(payload)
                temporary.replace(target)
                original = self.media_dir / f"{key}.carrier.png"
                original.write_bytes(source)
                record.update(extension=extension, size=len(payload), url="/media/" + name,
                              carrier_url="/media/" + original.name)
                files[name] = target
                files[original.name] = original
            except (OSError, ValueError) as exc:
                # A partially written PNG can be retried after its size or mtime changes.
                record["error"] = str(exc)
            records.append(record)
        return records, files

    async def scan(self):
        async with self.scan_lock:
            try:
                records, files = await asyncio.to_thread(self._scan)
            except OSError:
                return False
            for record in records:
                self.media[record["id"]] = record
            self.files.update(files)
            if records:
                self.revision += 1
            return bool(records)

    async def publish(self):
        async with self.publish_lock:
            snapshot = self.snapshot()
            async def send(client):
                try:
                    await asyncio.wait_for(client.send_json(snapshot), timeout=2)
                except (OSError, RuntimeError, asyncio.TimeoutError):
                    self.clients.discard(client)
            await asyncio.gather(*(send(client) for client in tuple(self.clients)))


def create_app(state=None, scan_interval=1.0):
    state = state or ReceiverState()
    app = web.Application(client_max_size=1024 * 1024)

    async def home(request):
        return web.Response(text=PAGE, content_type="text/html", headers={"Cache-Control": "no-store"})

    async def script(request):
        return web.FileResponse(Path(__file__).with_name("local_results.js"), headers={"Cache-Control": "no-store"})

    async def receive_callback(request):
        if request.content_type != "application/json":
            raise web.HTTPUnsupportedMediaType(text="use application/json")
        try:
            body = json.loads((await request.read()).decode("utf-8"))
        except (UnicodeError, ValueError):
            raise web.HTTPBadRequest(text="request body must be UTF-8 JSON")
        if (not isinstance(body, dict) or set(body) not in ({"msg"}, {"id", "msg"})
                or not isinstance(body["msg"], str) or not isinstance(body.get("id", ""), str)):
            raise web.HTTPBadRequest(text='body must be {"id": string, "msg": string}; id is optional')
        try:
            state.receive(body["msg"], body.get("id", ""))
        except OSError:
            raise web.HTTPInternalServerError(text="could not save the message log")
        await state.publish()
        return web.json_response({"ok": True})

    async def callback(request):
        # Keep the requested test status even if validation or storage fails.
        try:
            response = await receive_callback(request)
        except web.HTTPException as exc:
            response = web.Response(text=exc.text, content_type="text/plain")
        except Exception:
            logging.exception("Local callback processing failed")
            response = web.json_response({"ok": False, "error": "callback processing failed"})
        response.set_status(200)
        return response

    async def events(request):
        if await state.scan():
            await state.publish()
        return web.json_response(state.snapshot(), headers={"Cache-Control": "no-store"})

    async def health(request):
        return web.json_response({"ok": True, "service": "chicktools-callback-receiver", "received": state.received})

    async def media(request):
        target = state.files.get(request.match_info["name"])
        if target is None or not target.is_file():
            raise web.HTTPNotFound()
        # Range requests allow seeking without reading the whole video into memory.
        return web.FileResponse(target, headers={"X-Content-Type-Options": "nosniff"})

    async def websocket(request):
        origin = request.headers.get("Origin")
        if origin and origin != f"{request.scheme}://{request.host}":
            raise web.HTTPForbidden()
        ws = web.WebSocketResponse(heartbeat=20, max_msg_size=4096)
        await ws.prepare(request)
        state.clients.add(ws)
        try:
            await ws.send_json(state.snapshot())
            async for _ in ws:
                pass
        finally:
            state.clients.discard(ws)
        return ws

    async def watch():
        while True:
            await asyncio.sleep(scan_interval)
            if await state.scan():
                await state.publish()

    async def lifecycle(app):
        await state.scan()
        watcher = asyncio.create_task(watch())
        yield
        watcher.cancel()
        with suppress(asyncio.CancelledError):
            await watcher

    async def shutdown(app):
        await asyncio.gather(*(ws.close(code=1001) for ws in tuple(state.clients)))

    app.cleanup_ctx.append(lifecycle)
    app.on_shutdown.append(shutdown)
    app.add_routes([web.get("/", home), web.get("/viewer.js", script),
                    web.post("/callback", callback), web.get("/events", events),
                    web.get("/health", health), web.get("/ws", websocket), web.get("/media/{name}", media)])
    return app


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", default="127.0.0.1", choices=["127.0.0.1"])
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument("--log", type=Path, default=Path(".local/callback-received.jsonl"))
    parser.add_argument("--carrier-dir", type=Path, help="directory containing ChickSaveImage carrier PNGs")
    parser.add_argument("--password", default="", help="local demo carrier password")
    parser.add_argument("--media-dir", type=Path, default=Path(".local/decoded-media"))
    args = parser.parse_args()
    state = ReceiverState(args.log, args.carrier_dir, args.media_dir, args.password)
    print(f"Result viewer: http://{args.host}:{args.port}/", flush=True)
    web.run_app(create_app(state), host=args.host, port=args.port, access_log=None)


if __name__ == "__main__":
    main()