from __future__ import annotations

import argparse
import hashlib
import json
import os
from datetime import UTC, datetime
from pathlib import Path, PurePosixPath
from typing import Any

from cvat_sdk import make_client
from import_e2_workspace import _ensure_task, _label_specs, _sha256

WORKSPACE_SCHEMA = "missioncore.lab-e3-cvat-review-workspace/v1"
IDENTITY_SCHEMA = "missioncore.lab-e3-cvat-review-identity/v1"
EXPECTED_PACK_ID = (
    "evaluation-pack-"
    "7a983bba75d46c7c260252cb2d461e1384dcb92cda9e164397e841e6ebb37789"
)
EXPECTED_RESULT_ID = (
    "e3-segmentation-"
    "01bd497c44c2b940add145ec784d3418010327bce0baddd4420b0925317e8a16"
)
EXPECTED_FRAMES = 64


def _canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode()


def _valid_sha256(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) == 64
        and all(character in "0123456789abcdef" for character in value)
    )


def _read_object(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise RuntimeError(f"JSON root is not an object: {path}")
    return value


def _safe_artifact(root: Path, encoded: object) -> Path:
    if not isinstance(encoded, str):
        raise RuntimeError("artifact path is not a string")
    relative = PurePosixPath(encoded)
    if relative.is_absolute() or ".." in relative.parts or not relative.parts:
        raise RuntimeError("artifact path is unsafe")
    path = root.joinpath(*relative.parts).resolve(strict=True)
    if not path.is_file() or not path.is_relative_to(root):
        raise RuntimeError("artifact path escaped its root")
    return path


def _workspace(root: Path, images_zip: Path) -> tuple[dict[str, Any], dict[str, Path]]:
    manifest_path = root / "manifest.json"
    manifest = _read_object(manifest_path)
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    if (
        manifest.get("schema_version") != WORKSPACE_SCHEMA
        or manifest.get("ground_truth") is not False
        or not isinstance(identity, dict)
        or identity.get("schema_version") != IDENTITY_SCHEMA
        or not _valid_sha256(identity_sha256)
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or manifest.get("workspace_id") != f"e3-cvat-review-{identity_sha256}"
        or root.name != manifest.get("workspace_id")
        or identity.get("evaluation_pack_id") != EXPECTED_PACK_ID
        or identity.get("e3_result_id") != EXPECTED_RESULT_ID
        or identity.get("frame_count") != EXPECTED_FRAMES
        or _sha256(images_zip) != identity.get("images_zip_sha256")
    ):
        raise RuntimeError("LAB E3 CVAT review workspace is incompatible")
    artifacts: dict[str, Path] = {}
    descriptors = manifest.get("artifacts")
    if not isinstance(descriptors, list):
        raise RuntimeError("LAB E3 CVAT artifact list is invalid")
    for descriptor in descriptors:
        if not isinstance(descriptor, dict):
            raise RuntimeError("LAB E3 CVAT artifact descriptor is invalid")
        path = _safe_artifact(root, descriptor.get("path"))
        if (
            path.stat().st_size != descriptor.get("bytes")
            or not _valid_sha256(descriptor.get("sha256"))
            or _sha256(path) != descriptor["sha256"]
        ):
            raise RuntimeError(f"LAB E3 CVAT artifact changed: {path.name}")
        artifacts[path.name] = path
    if set(artifacts) != {
        "labels.json",
        "control-fisheye-mask.zip",
        "challenger-kb4-cubemap5.zip",
    }:
        raise RuntimeError("LAB E3 CVAT artifact set changed")
    return manifest, artifacts


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", required=True)
    parser.add_argument("--username", required=True)
    parser.add_argument("--password-env", required=True)
    parser.add_argument("--workspace", type=Path, required=True)
    parser.add_argument("--images-zip", type=Path, required=True)
    parser.add_argument("--report", type=Path, required=True)
    args = parser.parse_args()

    password = os.environ.get(args.password_env)
    if not password:
        raise RuntimeError(f"Password environment variable {args.password_env!r} is empty")
    workspace = args.workspace.resolve(strict=True)
    images_zip = args.images_zip.resolve(strict=True)
    manifest, artifacts = _workspace(workspace, images_zip)
    labels_document = _read_object(artifacts["labels.json"])
    semantic_labels = labels_document.get("semantic_task")
    if not isinstance(semantic_labels, list):
        raise RuntimeError("semantic labels are absent")
    labels = _label_specs(semantic_labels)
    task_specs = (
        {
            "name": "LAB E3 | K1 | EoMT fisheye control | pack 7a983bba",
            "annotation": artifacts["control-fisheye-mask.zip"],
            "role": "control",
        },
        {
            "name": "LAB E3 | K1 | EoMT KB4 cubemap5 challenger | pack 7a983bba",
            "annotation": artifacts["challenger-kb4-cubemap5.zip"],
            "role": "challenger",
        },
    )
    with make_client(args.server, credentials=(args.username, password)) as client:
        client.check_server_version(fail_if_unsupported=True)
        tasks = []
        for spec in task_specs:
            record = _ensure_task(
                client,
                name=spec["name"],
                labels=labels,
                images_path=images_zip,
                annotation_path=spec["annotation"],
                annotation_format="Segmentation mask 1.1",
                expected_shape_count=None,
            )
            record["role"] = spec["role"]
            tasks.append(record)

    report = {
        "schema_version": "missioncore.lab-e3-cvat-import/v1",
        "created_at_utc": datetime.now(UTC)
        .isoformat(timespec="milliseconds")
        .replace("+00:00", "Z"),
        "server": args.server,
        "cvat_version": "v2.70.0",
        "workspace_id": manifest["workspace_id"],
        "workspace_manifest_sha256": _sha256(workspace / "manifest.json"),
        "evaluation_pack_id": EXPECTED_PACK_ID,
        "e3_result_id": EXPECTED_RESULT_ID,
        "ground_truth": False,
        "priority_image_ids": manifest["identity"]["priority_image_ids"],
        "tasks": tasks,
        "next_gate": "two-pass human review and reviewed export",
    }
    args.report.parent.mkdir(parents=True, exist_ok=True)
    args.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()
