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

from __future__ import annotations

import asyncio
import os
from datetime import datetime
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.artifact_lifecycle import register_published_artifacts
from core.paths import from_db_path
from core.software_contracts import (
    SoftwareContractError,
    default_capabilities,
    get_contract,
)
from core.software_jobs import (
    MAX_OUTPUT_ARTIFACT_BYTES,
    MAX_OUTPUT_TOTAL_BYTES,
    SoftwareJobError,
    abandon_offer,
    create_job,
    get_job,
    get_job_input,
    get_job_output_context,
    list_jobs,
    mark_node_jobs_disconnected,
    offer_next_job,
    pending_node_cancellations,
    record_job_terminal,
    replay_succeeded_outputs,
    request_job_analysis,
    request_job_cancel,
    respond_to_offer,
    software_job_output_path,
    succeeded_output_upload_matches,
    update_job_state,
    validate_output_manifest,
)
from core.software_nodes import (
    SoftwareNodeError,
    authenticate_node,
    create_enrollment,
    delete_node,
    enroll_node,
    list_nodes,
    mark_node_offline,
    set_node_disabled,
    update_node_runtime,
)
from web.schemas import (
    SoftwareEnrollmentCreateRequest,
    SoftwareJobCreateRequest,
    SoftwareNodeDisableRequest,
    SoftwareNodeEnrollRequest,
)
from web.software_followups import dispatch_followup
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 SoftwareNodeError("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, SoftwareNodeError) 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, "software 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, "software job output path contains a symbolic link")


def _task_working_dir(root: Path, stored: str) -> Path:
    """解析 tasks.working_dir 的 ROOT 相对 DB 形态，并再次守住 user_root 边界。"""
    working_dir = from_db_path(stored).resolve()
    try:
        working_dir.relative_to(root.resolve())
    except ValueError as exc:
        raise SoftwareJobError("software job working directory is outside user root") from exc
    return working_dir


def _staged_output_path(
    staging: Path, capability: str, output_id: str, filename: str
) -> Path:
    flat = staging / filename
    organized = staging / software_job_output_path(capability, output_id)
    return flat if flat.is_file() else organized


def _organize_staged_metadata(
    staging: Path, capability: str, manifest: list[dict]
) -> None:
    contract = get_contract(capability)
    metadata = [
        item for item in manifest
        if not contract.output_spec(item["artifact_id"]).publish
    ]
    if not metadata:
        return
    (staging / ".meta").mkdir(exist_ok=True)
    for item in metadata:
        source = staging / item["filename"]
        destination = staging / software_job_output_path(capability, item["artifact_id"])
        if not source.is_file():
            if destination.is_file():
                continue
            raise SoftwareJobError(f"uploaded artifact is missing: {item['artifact_id']}")
        if destination.exists():
            if (
                destination.is_file()
                and destination.stat().st_size == item["size_bytes"]
                and _hash_file(destination) == item["sha256"]
            ):
                source.unlink()
                continue
            raise SoftwareJobError("software job metadata destination conflicts")
        os.replace(source, destination)


def _publish_software_job_outputs(job_id: UUID, context: dict, manifest: list[dict]) -> list[dict]:
    capability = context.get("capability", default_capabilities()[0])
    contract = get_contract(capability)
    root = load_user_root(context["user_id"])
    working_dir = _task_working_dir(root, context["working_dir"])
    staging = safe_join(root, f".zcbot_software_job_staging/{job_id}")
    relative_output = Path(contract.output_namespace) / 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 = (
            _staged_output_path(
                staging, capability, item["artifact_id"], item["filename"]
            )
            if source == staging
            else destination / software_job_output_path(capability, item["artifact_id"])
        )
        if (
            not path.is_file()
            or path.stat().st_size != item["size_bytes"]
            or _hash_file(path) != item["sha256"]
        ):
            raise SoftwareJobError(f"uploaded artifact is missing or invalid: {item['artifact_id']}")
    if source == staging:
        _organize_staged_metadata(staging, capability, manifest)
        destination.parent.mkdir(parents=True, exist_ok=True)
        if destination.exists():
            raise SoftwareJobError("software job output destination already exists unexpectedly")
        os.replace(staging, destination)
        try:
            staging.parent.rmdir()
        except OSError:
            pass
    refs = tuple({
        "path": (
            relative_output / software_job_output_path(capability, item["artifact_id"])
        ).as_posix(),
        "label": item["filename"],
        "media_type": item["media_type"],
    } for item in manifest if contract.output_spec(item["artifact_id"]).publish)
    published_refs = register_published_artifacts(
        user_id=context["user_id"],
        task_id=context["task_id"],
        user_root=root,
        working_dir=working_dir,
        refs=refs,
        software_job_id=job_id,
    )
    refs_by_path = {item["path"]: item for item in published_refs}
    return [
        {
            **item,
            "source_artifact_id": item["artifact_id"],
            "artifact_id": (
                refs_by_path.get(
                    (
                        relative_output
                        / software_job_output_path(capability, item["artifact_id"])
                    ).as_posix(),
                    {},
                ).get("artifact_id")
            ),
            "path": (
                relative_output / software_job_output_path(capability, item["artifact_id"])
            ).as_posix(),
        }
        for item in manifest
    ]


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

    @app.get("/v1/software-jobs/{job_id}/inputs/{input_key}", tags=["software-nodes"])
    def download_software_job_input(
        job_id: UUID,
        input_key: str,
        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, SoftwareNodeError) as exc:
            raise HTTPException(401, "invalid node credentials") from exc
        item = get_job_input(node_id, job_id, input_key)
        if item is None:
            raise HTTPException(404, "software job input not found")
        target = safe_join(load_user_root(item["user_id"]), item["current_path"])
        if not target.is_file():
            raise HTTPException(404, "software job input file not found")
        stat = target.stat()
        if stat.st_size != item["size_bytes"]:
            raise HTTPException(409, "software 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, "software 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/software-jobs/{job_id}/outputs/{artifact_id}",
        tags=["software-nodes"],
        status_code=status.HTTP_204_NO_CONTENT,
    )
    async def upload_software_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,
        )
        try:
            contract = get_contract(context["capability"])
            output_spec = contract.output_spec(artifact_id)
            requested_ids = set(contract.expected_outputs(context["request"]))
        except SoftwareContractError as exc:
            raise HTTPException(400, str(exc)) from exc
        if artifact_id not in requested_ids:
            raise HTTPException(400, "unsupported output artifact identity")
        filename = output_spec.filename
        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")
        try:
            if succeeded_output_upload_matches(
                context,
                artifact_id,
                size_bytes=x_content_length,
                digest=x_content_sha256,
            ):
                return
        except SoftwareJobError as exc:
            raise HTTPException(409, str(exc)) from exc
        root = load_user_root(context["user_id"])
        working_dir = _task_working_dir(root, context["working_dir"])
        published = safe_join(
            working_dir,
            (
                f"{contract.output_namespace}/{job_id}/"
                f"{software_job_output_path(context['capability'], artifact_id)}"
            ),
        )
        if published.is_file():
            if published.stat().st_size == x_content_length and _hash_file(published) == x_content_sha256:
                return
            raise HTTPException(409, "published output conflicts with uploaded artifact")
        staging = safe_join(root, f".zcbot_software_job_staging/{job_id}")
        _reject_symlink_path(root, staging)
        staging.mkdir(parents=True, exist_ok=True)
        organized = staging / software_job_output_path(context["capability"], artifact_id)
        if organized.is_file():
            if organized.stat().st_size == x_content_length and _hash_file(organized) == x_content_sha256:
                return
            raise HTTPException(409, "staged output conflicts with uploaded artifact")
        destination = staging / filename
        if destination.is_file():
            if destination.stat().st_size == x_content_length and _hash_file(destination) == x_content_sha256:
                return
            raise HTTPException(409, "uploaded output conflicts with existing staging file")
        staged_total = sum(item.stat().st_size for item in staging.rglob("*") if item.is_file())
        if staged_total + x_content_length > MAX_OUTPUT_TOTAL_BYTES:
            raise HTTPException(413, "software 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

    @app.post("/v1/software-jobs/{job_id}/outputs/complete", tags=["software-nodes"])
    async def complete_software_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["capability"], context["request"], body.get("artifact_manifest")
            )
            replayed = replay_succeeded_outputs(context, manifest)
            if replayed is not None:
                return {"status": "succeeded", "artifact_manifest": replayed}
            published = await asyncio.to_thread(
                _publish_software_job_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)
            await dispatch_followup(request.app, job_id)
        except (SoftwareJobError, KeyError, TypeError) as exc:
            raise HTTPException(409, str(exc)) from exc
        return {"status": "succeeded", "artifact_manifest": published}

    @app.websocket("/v1/software-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, SoftwareNodeError):
            # 握手前 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},
            )
            for cancel in await asyncio.to_thread(pending_node_cancellations, node_id):
                await node_connections.send_on(
                    node_id, websocket, {"type": "job_cancel", "payload": cancel}
                )
            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")}
                )
                for cancel in await asyncio.to_thread(pending_node_cancellations, node_id):
                    await node_connections.send_on(
                        node_id, websocket, {"type": "job_cancel", "payload": cancel}
                    )
                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 (
            SoftwareJobError,
            SoftwareNodeError,
            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}/software-jobs", tags=["software-jobs"])
    async def submit_software_job(
        task_id: UUID,
        body: SoftwareJobCreateRequest,
        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 SoftwareJobError as exc:
            detail = str(exc)
            raise HTTPException(404 if detail == "task not found" else 400, detail) from exc

    @app.get("/v1/software-jobs", tags=["software-jobs"])
    def read_software_jobs(
        task_id: UUID | None = None,
        active_only: bool = False,
        limit: int = 50,
        before_created_at: datetime | None = None,
        before_job_id: UUID | None = None,
        user_id: UUID = Depends(require_user),  # noqa: B008
    ):
        if (before_created_at is None) != (before_job_id is None):
            raise HTTPException(
                400, "before_created_at and before_job_id must be provided together"
            )
        page_limit = max(1, min(int(limit), 100))
        before = (
            (before_created_at, before_job_id)
            if before_created_at is not None and before_job_id is not None
            else None
        )
        results = list_jobs(
            user_id, task_id=task_id, active_only=active_only,
            limit=page_limit + 1, before=before,
        )
        has_more = len(results) > page_limit
        results = results[:page_limit]
        last = results[-1] if has_more else None
        return {
            "results": results,
            "next_cursor": (
                {"created_at": last["created_at"], "job_id": last["job_id"]}
                if last is not None else None
            ),
        }

    @app.post("/v1/software-jobs/{job_id}/cancel", tags=["software-jobs"])
    async def cancel_software_job(
        job_id: UUID,
        user_id: UUID = Depends(require_user),  # noqa: B008
    ):
        try:
            job, message = await asyncio.to_thread(request_job_cancel, user_id, job_id)
        except SoftwareJobError as exc:
            raise HTTPException(404, str(exc)) from exc
        if message is not None:
            await node_connections.send(
                message["node_id"],
                {"type": "job_cancel", "payload": message["payload"]},
            )
        return job

    @app.post("/v1/software-jobs/{job_id}/analyze", tags=["software-jobs"])
    async def analyze_software_job(
        job_id: UUID,
        request: Request,
        user_id: UUID = Depends(require_user),  # noqa: B008
    ):
        try:
            await asyncio.to_thread(request_job_analysis, user_id, job_id)
        except SoftwareJobError as exc:
            detail = str(exc)
            raise HTTPException(404 if detail == "job not found" else 409, detail) from exc
        await dispatch_followup(request.app, job_id)
        return get_job(user_id, job_id)

    @app.get("/v1/software-jobs/{job_id}", tags=["software-jobs"])
    def read_software_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, "software job not found")
        return job

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

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

    @app.patch("/v1/admin/software-nodes/{node_id}", tags=["admin"])
    async def admin_disable_software_node(
        node_id: UUID,
        body: SoftwareNodeDisableRequest,
        user_id: UUID = Depends(require_admin),  # noqa: B008
    ):
        if not await asyncio.to_thread(set_node_disabled, node_id, body.disabled):
            raise HTTPException(404, "software 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/software-nodes/{node_id}", tags=["admin"])
    async def admin_delete_software_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, "software node not found")
        return {"node_id": str(node_id), "status": "deleted"}