"""Receive original AUVRA media and allow the canvas to use ComfyUI."""

import asyncio
from collections import Counter
import hashlib
import http.client
import ipaddress
import json
import os
import re
import shutil
import socket
import tempfile
import threading
import time
import urllib.parse
import urllib.request
from pathlib import Path

import folder_paths
from aiohttp import web
from server import PromptServer


ALLOWED_ORIGINS = {
    "https://auvra.art",
    "https://www.auvra.art",
    "http://127.0.0.1:4175",
    "http://localhost:4175",
}


MAX_MEDIA_BYTES = 50 * 1024 * 1024
MAX_DIRECTORY_BYTES = 10 * 1024 * 1024 * 1024
IMPORT_TTL = 24 * 60 * 60
IMPORT_LOCK = threading.RLock()
ACTIVE_IMPORTS = Counter()
RESERVED_BYTES = 0
IMPORT_SLOTS = None
CLEANUP_TASK = None


def configured_cos_host():
    try:
        host = json.loads(Path(__file__).with_name("media-import.json").read_text("utf-8"))["cosHost"]
        if isinstance(host, str) and re.fullmatch(r"[a-z0-9-]+-\d+\.cos\.[a-z0-9-]+\.myqcloud\.com", host):
            return host
    except (OSError, ValueError, KeyError, TypeError):
        pass
    return None


COS_HOST = configured_cos_host()


def validate_source(url, export=False):
    parsed = urllib.parse.urlsplit(url)
    name = r"comfy-out-[A-Za-z0-9_.-]{1,170}" if export else r"comfy-[A-Za-z0-9_.-]{1,174}"
    if (parsed.scheme != "https" or parsed.username is not None or parsed.password is not None
            or parsed.port not in (None, 443) or parsed.fragment
            or not COS_HOST or parsed.hostname != COS_HOST
            or not re.fullmatch(r"/auvra/comfy-transfer/[a-f0-9]{16}/" + name, parsed.path)):
        raise ValueError("Only signed AUVRA COS media is accepted")
    query = urllib.parse.parse_qs(parsed.query, strict_parsing=True)
    if (any(len(values) != 1 for values in query.values())
            or query.get("q-sign-algorithm") != ["sha1"]
            or not re.fullmatch(r"[a-f0-9]{40}", query.get("q-signature", [""])[0])):
        raise ValueError("Invalid COS signature")
    start, end = map(int, query.get("q-key-time", [""])[0].split(";"))
    if not start <= time.time() + 60 or not time.time() < end <= time.time() + 900 or end - start > 900:
        raise ValueError("Expired or excessive COS access duration")
    addresses = {item[4][0] for item in socket.getaddrinfo(parsed.hostname, 443, type=socket.SOCK_STREAM)}
    if not addresses or any(not ipaddress.ip_address(value).is_global for value in addresses):
        raise ValueError("Non-public source address")
    return min(addresses, key=lambda value: (":" in value, value))


class PinnedConnection(http.client.HTTPSConnection):
    def __init__(self, host, *, pinned_ip, **kwargs):
        self.pinned_ip = pinned_ip
        super().__init__(host, **kwargs)

    def connect(self):
        self.sock = self._create_connection((self.pinned_ip, self.port), self.timeout, self.source_address)
        self.sock = self._context.wrap_socket(self.sock, server_hostname=self.host)


class PinnedHandler(urllib.request.HTTPSHandler):
    def __init__(self, pinned_ip):
        super().__init__()
        self.pinned_ip = pinned_ip

    def https_open(self, request):
        return self.do_open(lambda host, **kwargs: PinnedConnection(host, pinned_ip=self.pinned_ip, **kwargs), request)


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None


def queued_import_names():
    running, pending = PromptServer.instance.prompt_queue.get_current_queue()
    names = set()
    def collect(value):
        if isinstance(value, str) and value.startswith("auvra/"):
            names.add(value[6:])
        elif isinstance(value, dict):
            for child in value.values():
                collect(child)
        elif isinstance(value, (list, tuple)):
            for child in value:
                collect(child)
    collect((running, pending))
    return names


def cleanup_imports(directory=None, required=0):
    directory = directory or Path(folder_paths.get_input_directory()) / "auvra"
    with IMPORT_LOCK:
        if not directory.exists():
            return
        try:
            protected = queued_import_names() | set(ACTIVE_IMPORTS)
        except Exception:
            # A queue inspection failure cannot justify deleting a job's inputs.
            protected = None
        now, total, files = time.time(), 0, []
        for path in directory.iterdir():
            if path.is_symlink() or not path.is_file():
                continue
            stat = path.stat()
            if path.name.startswith(".import-") and ACTIVE_IMPORTS[path.name]:
                continue  # Its full size is already reserved before downloading.
            total += stat.st_size
            owned = re.fullmatch(r"[a-f0-9]{64}\.[a-z0-9]{2,5}|\.import-[A-Za-z0-9_-]+", path.name)
            if owned and protected is not None and path.name not in protected:
                files.append((stat.st_mtime, path, stat.st_size))
        for modified, path, size in sorted(files):
            expired = now - modified > IMPORT_TTL
            over_limit = total + RESERVED_BYTES + required > MAX_DIRECTORY_BYTES
            # Allow the caller to finish preparing references before /prompt queues them.
            if expired or (over_limit and now - modified > 600):
                path.unlink()
                total -= size
        if total + RESERVED_BYTES + required > MAX_DIRECTORY_BYTES:
            raise OSError("Original media capacity reached; active/recent inputs are retained")


async def cleanup_loop():
    while True:
        try:
            await asyncio.get_running_loop().run_in_executor(None, cleanup_imports)
        except Exception as exc:
            print(f"[AUVRA Bridge] input cleanup deferred: {type(exc).__name__}")
        await asyncio.sleep(600)


async def start_cleanup(_app):
    global CLEANUP_TASK
    CLEANUP_TASK = asyncio.create_task(cleanup_loop())


async def stop_cleanup(_app):
    if CLEANUP_TASK is not None:
        CLEANUP_TASK.cancel()
        try:
            await CLEANUP_TASK
        except asyncio.CancelledError:
            pass


def import_original(payload):
    global RESERVED_BYTES
    size, digest, suffix = payload.get("size"), payload.get("sha256"), payload.get("suffix")
    if (type(size) is not int or not 0 < size <= MAX_MEDIA_BYTES
            or not isinstance(digest, str) or not re.fullmatch(r"[a-f0-9]{64}", digest)
            or not isinstance(suffix, str) or suffix not in {".png", ".jpg", ".jpeg", ".jpe", ".webp", ".gif", ".bmp", ".tif", ".tiff", ".avif", ".mp4", ".webm", ".mov", ".mkv", ".wav", ".mp3", ".m4a", ".flac", ".ogg", ".aac"}):
        raise ValueError("Invalid media identity or size")
    url = payload.get("url")
    if not isinstance(url, str) or len(url) > 4096:
        raise ValueError("Invalid media source")
    pinned_ip = validate_source(url)
    directory = Path(folder_paths.get_input_directory()) / "auvra"
    directory.mkdir(parents=True, exist_ok=True)
    filename = f"{digest}{suffix}"
    target = directory / filename
    result = {"name": filename, "subfolder": "auvra", "type": "input", "size": size, "sha256": digest}
    with IMPORT_LOCK:
        if target.is_symlink():
            raise ValueError("Invalid original media target")
        if target.is_file() and target.stat().st_size == size and hashlib.sha256(target.read_bytes()).hexdigest() == digest:
            target.touch()
            cleanup_imports(directory)
            return result
        cleanup_imports(directory, required=size)
        if shutil.disk_usage(directory).free < RESERVED_BYTES + size + 16 * 1024 * 1024:
            raise OSError("Insufficient disk space for original media")
        RESERVED_BYTES += size
        ACTIVE_IMPORTS[filename] += 1
    temporary = None
    try:
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect(), PinnedHandler(pinned_ip))
        started, count, checksum = time.monotonic(), 0, hashlib.sha256()
        with opener.open(url, timeout=20) as response, tempfile.NamedTemporaryFile(dir=directory, prefix=".import-", delete=False) as output:
            temporary = Path(output.name)
            with IMPORT_LOCK:
                ACTIVE_IMPORTS[temporary.name] += 1
            if response.status != 200 or (response.headers.get("Content-Length") and int(response.headers["Content-Length"]) != size):
                raise ValueError("Unexpected media response size or status")
            while True:
                if time.monotonic() - started > 120:
                    raise TimeoutError("Original media download timed out")
                chunk = response.read1(min(65536, size - count + 1))
                if not chunk:
                    break
                count += len(chunk)
                if count > size:
                    raise ValueError("Media exceeds declared size")
                checksum.update(chunk)
                output.write(chunk)
        if count != size or checksum.hexdigest() != digest:
            raise ValueError("Original media checksum mismatch")
        temporary.chmod(0o640)
        with IMPORT_LOCK:
            os.replace(temporary, target)
        return result
    finally:
        with IMPORT_LOCK:
            RESERVED_BYTES -= size
            ACTIVE_IMPORTS[filename] -= 1
            if not ACTIVE_IMPORTS[filename]:
                del ACTIVE_IMPORTS[filename]
            if temporary is not None:
                ACTIVE_IMPORTS[temporary.name] -= 1
                if not ACTIVE_IMPORTS[temporary.name]:
                    del ACTIVE_IMPORTS[temporary.name]
                temporary.unlink(missing_ok=True)


@PromptServer.instance.routes.get("/auvra/media-import")
async def import_capabilities(_request):
    return web.json_response({"originalImport": 1 if COS_HOST else 0, "originalExport": 1 if COS_HOST else 0, "maxBytes": MAX_MEDIA_BYTES, "maxExportBytes": MAX_EXPORT_BYTES})


@PromptServer.instance.routes.post("/auvra/media-import")
async def import_media(request):
    global IMPORT_SLOTS
    if IMPORT_SLOTS is None:
        IMPORT_SLOTS = asyncio.Semaphore(2)
    if request.content_length is not None and request.content_length > 4096:
        raise web.HTTPRequestEntityTooLarge(max_size=4096, actual_size=request.content_length)
    body = bytearray()
    while len(body) <= 4096:
        chunk = await request.content.read(4097 - len(body))
        if not chunk:
            break
        body.extend(chunk)
    if len(body) > 4096:
        raise web.HTTPRequestEntityTooLarge(max_size=4096, actual_size=len(body))
    try:
        payload = json.loads(body)
        if not isinstance(payload, dict):
            raise ValueError("Invalid import request")
        await asyncio.wait_for(IMPORT_SLOTS.acquire(), timeout=5)
    except (ValueError, TypeError):
        return web.json_response({"error": "Invalid original media import request"}, status=400)
    except asyncio.TimeoutError:
        return web.json_response({"error": "Original media import busy"}, status=429)
    worker = asyncio.get_running_loop().run_in_executor(None, import_original, payload)
    try:
        return web.json_response(await asyncio.shield(worker))
    except Exception as exc:
        # Never return or log the signed URL.
        print(f"[AUVRA Bridge] original import failed: {type(exc).__name__}")
        return web.json_response({"error": f"Original media import failed ({type(exc).__name__})"}, status=400 if isinstance(exc, (ValueError, TypeError)) else 502)
    finally:
        if worker.done():
            IMPORT_SLOTS.release()
        else:
            def release(_worker):
                _worker.exception()
                IMPORT_SLOTS.release()
            worker.add_done_callback(release)


# Result export: the finished video goes from this GPU to AUVRA's transfer bucket, so AUVRA never has to pull large
# outputs back through this instance's public entrance (a slow cross-border link for a GPU in mainland China).
MAX_EXPORT_BYTES = 128 * 1024 * 1024
EXPORT_SLOTS = None


def export_source(payload):
    """The finished file inside ComfyUI's own output / temp directory; nothing else on this machine can be sent."""
    filename, subfolder, kind = payload.get("filename"), payload.get("subfolder") or "", payload.get("type") or "output"
    if (not isinstance(filename, str) or not re.fullmatch(r"[A-Za-z0-9_. ()-]{1,200}", filename) or filename.startswith(".")
            or not isinstance(subfolder, str) or len(subfolder) > 200 or ".." in subfolder.split("/") or subfolder.startswith("/")
            or kind not in {"output", "temp"}):
        raise ValueError("Invalid output reference")
    base = Path(folder_paths.get_output_directory() if kind == "output" else folder_paths.get_temp_directory()).resolve()
    path = (base / subfolder / filename)
    if path.is_symlink():
        raise ValueError("Invalid output reference")
    path = path.resolve()
    if base not in path.parents or not path.is_file():
        raise ValueError("Output file not found")
    size = path.stat().st_size
    if not 0 < size <= MAX_EXPORT_BYTES:
        raise ValueError("Output file is empty or too large")
    return path, size


def export_output(payload):
    url = payload.get("url")
    if not isinstance(url, str) or len(url) > 4096:
        raise ValueError("Invalid export target")
    pinned_ip = validate_source(url, export=True)
    path, size = export_source(payload)
    checksum = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            checksum.update(chunk)
    parsed = urllib.parse.urlsplit(url)
    connection = PinnedConnection(parsed.hostname, pinned_ip=pinned_ip, timeout=30)
    started = time.monotonic()
    try:
        connection.putrequest("PUT", f"{parsed.path}?{parsed.query}", skip_accept_encoding=True)
        connection.putheader("Content-Length", str(size))
        connection.putheader("Content-Type", "application/octet-stream")
        connection.endheaders()
        sent = 0
        with path.open("rb") as source:
            for chunk in iter(lambda: source.read(256 * 1024), b""):
                if time.monotonic() - started > 600:
                    raise TimeoutError("Output upload timed out")
                connection.send(chunk)
                sent += len(chunk)
        if sent != size:
            raise ValueError("Output changed while uploading")
        response = connection.getresponse()
        response.read(4096)
        if response.status != 200:
            raise ValueError(f"Transfer bucket refused the upload ({response.status})")
    finally:
        connection.close()
    return {"size": size, "sha256": checksum.hexdigest()}


@PromptServer.instance.routes.post("/auvra/media-export")
async def export_media(request):
    global EXPORT_SLOTS
    if EXPORT_SLOTS is None:
        EXPORT_SLOTS = asyncio.Semaphore(2)
    if request.content_length is not None and request.content_length > 8192:
        raise web.HTTPRequestEntityTooLarge(max_size=8192, actual_size=request.content_length)
    body = bytearray()
    while len(body) <= 8192:
        chunk = await request.content.read(8193 - len(body))
        if not chunk:
            break
        body.extend(chunk)
    if len(body) > 8192:
        raise web.HTTPRequestEntityTooLarge(max_size=8192, actual_size=len(body))
    try:
        payload = json.loads(body)
        if not isinstance(payload, dict):
            raise ValueError("Invalid export request")
        await asyncio.wait_for(EXPORT_SLOTS.acquire(), timeout=5)
    except (ValueError, TypeError):
        return web.json_response({"error": "Invalid output export request"}, status=400)
    except asyncio.TimeoutError:
        return web.json_response({"error": "Output export busy"}, status=429)
    worker = asyncio.get_running_loop().run_in_executor(None, export_output, payload)
    try:
        return web.json_response(await asyncio.shield(worker))
    except Exception as exc:
        # Never return or log the signed URL.
        print(f"[AUVRA Bridge] output export failed: {type(exc).__name__}")
        return web.json_response({"error": f"Output export failed ({type(exc).__name__})"}, status=400 if isinstance(exc, (ValueError, TypeError)) else 502)
    finally:
        if worker.done():
            EXPORT_SLOTS.release()
        else:
            def release(_worker):
                _worker.exception()
                EXPORT_SLOTS.release()
            worker.add_done_callback(release)


def _install():
    app = PromptServer.instance.app
    app.on_startup.append(start_cleanup)
    app.on_cleanup.append(stop_cleanup)
    for index, guard in enumerate(app.middlewares):
        if getattr(guard, "__name__", "") != "origin_only_middleware":
            continue

        @web.middleware
        async def auvra_local_origin(request, handler, _guard=guard):
            origin = request.headers.get("Origin", "")
            if origin not in ALLOWED_ORIGINS:
                return await _guard(request, handler)
            response = web.Response() if request.method == "OPTIONS" else await handler(request)
            if not response.prepared:
                response.headers["Access-Control-Allow-Origin"] = origin
                response.headers["Access-Control-Allow-Methods"] = "GET, POST, OPTIONS"
                response.headers["Access-Control-Allow-Headers"] = request.headers.get(
                    "Access-Control-Request-Headers", "Content-Type"
                )
                response.headers["Vary"] = "Origin"
                if request.headers.get("Access-Control-Request-Private-Network") == "true":
                    response.headers["Access-Control-Allow-Private-Network"] = "true"
            return response

        app.middlewares[index] = auvra_local_origin
        print("[AUVRA Bridge] enabled for AUVRA canvas origins")
        return
    raise RuntimeError("ComfyUI origin middleware was not found")


_install()
NODE_CLASS_MAPPINGS = {}