"""Windows Node MVP 的注册、管理与长连接端点。"""

from __future__ import annotations

import asyncio
import os
from hashlib import sha256
from pathlib import Path
from uuid import UUID

from fastapi import Depends, Header, HTTPException, Request, WebSocket, WebSocketDisconnect, status
from fastapi.responses import FileResponse

from core.compute_nodes import (
    ComputeNodeError,
    authenticate_node,
    create_enrollment,
    delete_node,
    enroll_node,
    list_nodes,
    mark_node_offline,
    set_node_disabled,
    update_node_runtime,
)
from core.compute_jobs import (
    MAX_OUTPUT_ARTIFACT_BYTES,
    MAX_OUTPUT_TOTAL_BYTES,
    OUTPUT_ARTIFACTS,
    abandon_offer,
    create_job,
    get_job,
    get_job_input,
    get_job_output_context,
    mark_node_jobs_disconnected,
    offer_next_job,
    record_job_terminal,
    respond_to_offer,
    update_job_state,
    validate_output_manifest,
)
from core.artifact_lifecycle import register_published_artifacts
from web.schemas import (
    ComputeEnrollmentCreateRequest,
    ComputeJobCreateRequest,
    ComputeNodeDisableRequest,
    ComputeNodeEnrollRequest,
)
from web.userfiles import load_user_root, safe_join


class NodeConnectionManager:
    def __init__(self) -> None:
        self._connections: dict[UUID, WebSocket] = {}
        self._send_locks: dict[UUID, asyncio.Lock] = {}
        self._lock = asyncio.Lock()

    async def activate(self, node_id: UUID, websocket: WebSocket) -> None:
        async with self._lock:
            old = self._connections.get(node_id)
            self._connections[node_id] = websocket
            self._send_locks.setdefault(node_id, asyncio.Lock())
        if old is not None and old is not websocket:
            await old.close(code=4001, reason="replaced by a newer connection")

    async def remove(self, node_id: UUID, websocket: WebSocket) -> bool:
        async with self._lock:
            if self._connections.get(node_id) is websocket:
                self._connections.pop(node_id, None)
                self._send_locks.pop(node_id, None)
                return True
        return False

    async def close(self, node_id: UUID) -> None:
        async with self._lock:
            websocket = self._connections.pop(node_id, None)
            self._send_locks.pop(node_id, None)
        if websocket is not None:
            await websocket.close(code=4003, reason="node disabled")

    async def node_ids(self) -> set[UUID]:
        async with self._lock:
            return set(self._connections)

    async def send(self, node_id: UUID, message: dict) -> bool:
        async with self._lock:
            websocket = self._connections.get(node_id)
            send_lock = self._send_locks.get(node_id)
        if websocket is None or send_lock is None:
            return False
        async with send_lock:
            await websocket.send_json(message)
        return True

    async def send_on(self, node_id: UUID, websocket: WebSocket, message: dict) -> bool:
        async with self._lock:
            current = self._connections.get(node_id)
            send_lock = self._send_locks.get(node_id)
        if current is not websocket or send_lock is None:
            return False
        async with send_lock:
            await websocket.send_json(message)
        return True


node_connections = NodeConnectionManager()


def _bearer(authorization: str | None) -> str:
    scheme, _, token = (authorization or "").partition(" ")
    if scheme.lower() != "bearer" or not token:
        raise ComputeNodeError("missing node bearer token")
    return token


def _authenticate_output_request(
    job_id: UUID,
    authorization: str | None,
    x_node_id: str,
    x_lease_id: str,
    x_request_digest: str,
) -> tuple[UUID, UUID, dict]:
    try:
        node_id = UUID(x_node_id)
        lease_id = UUID(x_lease_id)
        authenticate_node(node_id, _bearer(authorization))
    except (ValueError, ComputeNodeError) as exc:
        raise HTTPException(401, "invalid node credentials or job identity") from exc
    context = get_job_output_context(node_id, job_id, lease_id, x_request_digest)
    if context is None:
        raise HTTPException(404, "compute job output target not found")
    return node_id, lease_id, context


def _hash_file(path: Path) -> str:
    digest = sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _reject_symlink_path(root: Path, target: Path) -> None:
    root = root.resolve()
    current = root
    for part in target.relative_to(root).parts:
        current = current / part
        if current.is_symlink():
            raise HTTPException(409, "compute output path contains a symbolic link")


def _publish_compute_outputs(job_id: UUID, context: dict, manifest: list[dict]) -> list[dict]:
    root = load_user_root(context["user_id"])
    working_dir = safe_join(root, context["working_dir"])
    staging = safe_join(root, f".zcbot_compute_staging/{job_id}")
    relative_output = Path("origin") / str(job_id)
    destination = safe_join(working_dir, relative_output.as_posix())
    source = staging if staging.is_dir() else destination
    _reject_symlink_path(root, source)
    _reject_symlink_path(root, destination)
    for item in manifest:
        path = source / item["filename"]
        if (
            not path.is_file()
            or path.stat().st_size != item["size_bytes"]
            or _hash_file(path) != item["sha256"]
        ):
            raise ComputeNodeError(f"uploaded artifact is missing or invalid: {item['artifact_id']}")
    if source == staging:
        destination.parent.mkdir(parents=True, exist_ok=True)
        if destination.exists():
            raise ComputeNodeError("compute output destination already exists unexpectedly")
        os.replace(staging, destination)
        try:
            staging.parent.rmdir()
        except OSError:
            pass
    refs = tuple({
        "path": (relative_output / item["filename"]).as_posix(),
        "label": item["filename"],
        "media_type": item["media_type"],
    } for item in manifest)
    published_refs = register_published_artifacts(
        user_id=context["user_id"],
        task_id=context["task_id"],
        user_root=root,
        working_dir=working_dir,
        refs=refs,
    )
    refs_by_path = {item["path"]: item for item in published_refs}
    return [
        {
            **item,
            "source_artifact_id": item["artifact_id"],
            "artifact_id": refs_by_path[(relative_output / item["filename"]).as_posix()]["artifact_id"],
            "path": (relative_output / item["filename"]).as_posix(),
        }
        for item in manifest
    ]


def register_compute_node_routes(app, *, require_user, require_admin) -> None:
    @app.post(
        "/v1/compute/nodes/enroll",
        tags=["compute-nodes"],
        status_code=status.HTTP_201_CREATED,
    )
    def node_enroll(body: ComputeNodeEnrollRequest):
        try:
            return enroll_node(**body.model_dump())
        except ComputeNodeError as exc:
            raise HTTPException(400, str(exc)) from exc

    @app.get("/v1/compute/jobs/{job_id}/input", tags=["compute-nodes"])
    def download_compute_job_input(
        job_id: UUID,
        authorization: str | None = Header(default=None),
        x_node_id: str = Header(default=""),
    ):
        try:
            node_id = UUID(x_node_id)
            authenticate_node(node_id, _bearer(authorization))
        except (ValueError, ComputeNodeError) as exc:
            raise HTTPException(401, "invalid node credentials") from exc
        item = get_job_input(node_id, job_id)
        if item is None:
            raise HTTPException(404, "compute job input not found")
        target = safe_join(load_user_root(item["user_id"]), item["current_path"])
        if not target.is_file():
            raise HTTPException(404, "compute job input file not found")
        stat = target.stat()
        if stat.st_size != item["size_bytes"]:
            raise HTTPException(409, "compute job input changed after submission")
        digest = sha256()
        with target.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
        if digest.hexdigest() != item["sha256"]:
            raise HTTPException(409, "compute job input changed after submission")
        return FileResponse(
            path=str(target),
            filename=item["filename"],
            media_type="application/octet-stream",
            headers={
                "Cache-Control": "no-store",
                "X-Content-SHA256": item["sha256"],
            },
        )

    @app.put(
        "/v1/compute/jobs/{job_id}/outputs/{artifact_id}",
        tags=["compute-nodes"],
        status_code=status.HTTP_204_NO_CONTENT,
    )
    async def upload_compute_job_output(
        job_id: UUID,
        artifact_id: str,
        request: Request,
        authorization: str | None = Header(default=None),
        x_node_id: str = Header(default=""),
        x_lease_id: str = Header(default=""),
        x_request_digest: str = Header(default=""),
        x_content_sha256: str = Header(default=""),
        x_content_length: int = Header(default=-1),
    ):
        _, _, context = await asyncio.to_thread(
            _authenticate_output_request,
            job_id, authorization, x_node_id, x_lease_id, x_request_digest,
        )
        metadata = OUTPUT_ARTIFACTS.get(artifact_id)
        if metadata is None:
            raise HTTPException(400, "unsupported output artifact identity")
        filename, _, output_format = metadata
        requested_formats = set(context["request"].get("output", {}).get("formats") or [])
        if output_format is not None and output_format not in requested_formats:
            raise HTTPException(400, "output artifact was not requested")
        if not 1 <= x_content_length <= MAX_OUTPUT_ARTIFACT_BYTES:
            raise HTTPException(400, "output artifact size is invalid")
        if len(x_content_sha256) != 64 or any(c not in "0123456789abcdef" for c in x_content_sha256):
            raise HTTPException(400, "output artifact digest is invalid")
        root = load_user_root(context["user_id"])
        published = safe_join(
            safe_join(root, context["working_dir"]),
            f"origin/{job_id}/{filename}",
        )
        if published.is_file():
            if published.stat().st_size == x_content_length and _hash_file(published) == x_content_sha256:
                return None
            raise HTTPException(409, "published output conflicts with uploaded artifact")
        staging = safe_join(root, f".zcbot_compute_staging/{job_id}")
        _reject_symlink_path(root, staging)
        staging.mkdir(parents=True, exist_ok=True)
        destination = staging / filename
        if destination.is_file():
            if destination.stat().st_size == x_content_length and _hash_file(destination) == x_content_sha256:
                return None
            raise HTTPException(409, "uploaded output conflicts with existing staging file")
        staged_total = sum(
            item.stat().st_size for item in staging.iterdir() if item.is_file()
        )
        if staged_total + x_content_length > MAX_OUTPUT_TOTAL_BYTES:
            raise HTTPException(413, "compute job outputs exceed the total size limit")
        temporary = destination.with_name(destination.name + ".tmp-" + os.urandom(8).hex())
        digest = sha256()
        total = 0
        try:
            with temporary.open("xb") as handle:
                async for chunk in request.stream():
                    total += len(chunk)
                    if total > x_content_length or total > MAX_OUTPUT_ARTIFACT_BYTES:
                        raise HTTPException(413, "output artifact exceeded declared size")
                    digest.update(chunk)
                    handle.write(chunk)
                handle.flush()
                os.fsync(handle.fileno())
            if total != x_content_length or digest.hexdigest() != x_content_sha256:
                raise HTTPException(400, "output artifact did not match declared metadata")
            os.replace(temporary, destination)
        finally:
            temporary.unlink(missing_ok=True)
        return None

    @app.post("/v1/compute/jobs/{job_id}/outputs/complete", tags=["compute-nodes"])
    async def complete_compute_job_outputs(
        job_id: UUID,
        request: Request,
        authorization: str | None = Header(default=None),
        x_node_id: str = Header(default=""),
        x_lease_id: str = Header(default=""),
        x_request_digest: str = Header(default=""),
    ):
        node_id, lease_id, context = await asyncio.to_thread(
            _authenticate_output_request,
            job_id, authorization, x_node_id, x_lease_id, x_request_digest,
        )
        body = await request.json()
        if not isinstance(body, dict):
            raise HTTPException(400, "output completion body must be an object")
        try:
            manifest = validate_output_manifest(context["request"], body.get("artifact_manifest"))
            published = await asyncio.to_thread(_publish_compute_outputs, job_id, context, manifest)
            terminal = {
                "job_id": str(job_id),
                "lease_id": str(lease_id),
                "request_digest": x_request_digest,
                "status": "succeeded",
                "error": {},
                "artifact_manifest": published,
            }
            await asyncio.to_thread(record_job_terminal, node_id, terminal)
        except (ComputeNodeError, KeyError, TypeError) as exc:
            raise HTTPException(409, str(exc)) from exc
        return {"status": "succeeded", "artifact_manifest": published}

    @app.websocket("/v1/compute/nodes/connect")
    async def node_connect(websocket: WebSocket):
        try:
            node_id = UUID(websocket.headers.get("x-node-id", ""))
            token = _bearer(websocket.headers.get("authorization"))
            identity = await asyncio.to_thread(authenticate_node, node_id, token)
        except (ValueError, ComputeNodeError):
            # 握手前 close 会被 ASGI 统一表现为 HTTP 403，客户端无法区分
            # “凭据无效”和“代理/路由没有正确转发 WebSocket”。先升级再用
            # 应用关闭码给已持有 Node ID/Token 的节点返回明确诊断。
            await websocket.accept()
            await websocket.close(code=4003, reason="invalid node credentials")
            return
        await websocket.accept()
        await node_connections.activate(node_id, websocket)
        try:
            await node_connections.send_on(
                node_id,
                websocket,
                {"type": "connected", "heartbeat_seconds": 15},
            )
            while True:
                message = await websocket.receive_json()
                message_type = message.get("type")
                payload = message.get("payload") or {}
                if not isinstance(payload, dict):
                    await node_connections.send_on(node_id, websocket,
                        {"type": "error", "code": "unsupported_message"}
                    )
                    continue
                if message_type in {"job_accept", "job_reject"}:
                    await asyncio.to_thread(
                        respond_to_offer,
                        node_id,
                        accepted=message_type == "job_accept",
                        payload=payload,
                    )
                    await node_connections.send_on(node_id, websocket,
                        {"type": "ack", "message_id": message.get("message_id")}
                    )
                    continue
                if message_type in {"job_state", "job_terminal"}:
                    if message_type == "job_terminal" and payload.get("status") == "succeeded":
                        await node_connections.send_on(node_id, websocket,
                            {
                                "type": "error",
                                "code": "outputs_not_published",
                                "message_id": message.get("message_id"),
                            }
                        )
                        continue
                    handler = (
                        update_job_state
                        if message_type == "job_state"
                        else record_job_terminal
                    )
                    await asyncio.to_thread(handler, node_id, payload)
                    await node_connections.send_on(node_id, websocket,
                        {"type": "ack", "message_id": message.get("message_id")}
                    )
                    continue
                if message_type not in {"hello", "heartbeat"}:
                    await node_connections.send_on(node_id, websocket,
                        {"type": "error", "code": "unsupported_message"}
                    )
                    continue
                if payload.get("install_id") and payload["install_id"] != str(
                    identity["install_id"]
                ):
                    await websocket.close(code=1008, reason="install identity mismatch")
                    return
                await asyncio.to_thread(
                    update_node_runtime,
                    node_id,
                    status="online",
                    runtime=payload,
                )
                await node_connections.send_on(node_id, websocket,
                    {"type": "ack", "message_id": message.get("message_id")}
                )
                offer = await asyncio.to_thread(
                    offer_next_job, await node_connections.node_ids()
                )
                if offer is not None:
                    delivered = await node_connections.send(
                        offer["node_id"],
                        {"type": "job_offer", "payload": offer["payload"]},
                    )
                    if not delivered:
                        await asyncio.to_thread(
                            abandon_offer, offer["node_id"], offer["payload"]
                        )
        except (ComputeNodeError, WebSocketDisconnect, RuntimeError, ValueError):
            pass
        finally:
            if await node_connections.remove(node_id, websocket):
                await asyncio.to_thread(mark_node_offline, node_id)
                await asyncio.to_thread(mark_node_jobs_disconnected, node_id)

    @app.post("/v1/tasks/{task_id}/compute-jobs", tags=["compute-jobs"])
    async def submit_compute_job(
        task_id: UUID,
        body: ComputeJobCreateRequest,
        user_id: UUID = Depends(require_user),  # noqa: B008
    ):
        try:
            job, created = await asyncio.to_thread(
                create_job, user_id, task_id, **body.model_dump()
            )
            offer = await asyncio.to_thread(
                offer_next_job, await node_connections.node_ids()
            )
            if offer is not None:
                delivered = await node_connections.send(
                    offer["node_id"], {"type": "job_offer", "payload": offer["payload"]}
                )
                if not delivered:
                    await asyncio.to_thread(
                        abandon_offer, offer["node_id"], offer["payload"]
                    )
            return {**job, "created": created}
        except ComputeNodeError as exc:
            detail = str(exc)
            raise HTTPException(404 if detail == "task not found" else 400, detail) from exc

    @app.get("/v1/compute-jobs/{job_id}", tags=["compute-jobs"])
    def read_compute_job(
        job_id: UUID,
        user_id: UUID = Depends(require_user),  # noqa: B008
    ):
        job = get_job(user_id, job_id)
        if job is None:
            raise HTTPException(404, "compute job not found")
        return job

    @app.post("/v1/admin/compute-node-enrollments", tags=["admin"])
    def admin_create_compute_enrollment(
        body: ComputeEnrollmentCreateRequest,
        user_id: UUID = Depends(require_admin),  # noqa: B008
    ):
        try:
            return create_enrollment(user_id, **body.model_dump())
        except ComputeNodeError as exc:
            raise HTTPException(400, str(exc)) from exc

    @app.get("/v1/admin/compute-nodes", tags=["admin"])
    def admin_compute_nodes(user_id: UUID = Depends(require_admin)):  # noqa: B008
        return {"results": list_nodes()}

    @app.patch("/v1/admin/compute-nodes/{node_id}", tags=["admin"])
    async def admin_disable_compute_node(
        node_id: UUID,
        body: ComputeNodeDisableRequest,
        user_id: UUID = Depends(require_admin),  # noqa: B008
    ):
        if not await asyncio.to_thread(set_node_disabled, node_id, body.disabled):
            raise HTTPException(404, "compute node not found")
        if body.disabled:
            await node_connections.close(node_id)
        return {
            "node_id": str(node_id),
            "status": "disabled" if body.disabled else "offline",
        }

    @app.delete("/v1/admin/compute-nodes/{node_id}", tags=["admin"])
    async def admin_delete_compute_node(
        node_id: UUID,
        user_id: UUID = Depends(require_admin),  # noqa: B008
    ):
        # 先撤掉在线连接，避免删除后的旧 socket 继续上报运行态。
        await node_connections.close(node_id)
        if not await asyncio.to_thread(delete_node, node_id):
            raise HTTPException(404, "compute node not found")
        return {"node_id": str(node_id), "status": "deleted"}