from __future__ import annotations

import argparse
import hashlib
import json
import os
import time
from collections import Counter
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from cvat_sdk import make_client, models
from cvat_sdk.api_client.exceptions import ServiceException
from cvat_sdk.core.proxies.tasks import ResourceType

WORKSPACE_ID = (
    "annotation-workspace-"
    "9a950d1c37d56dc12cc285b13c5addd7795285879cbcb1fbb2d5811c3c69821a"
)
EVALUATION_PACK_ID = (
    "evaluation-pack-"
    "7a983bba75d46c7c260252cb2d461e1384dcb92cda9e164397e841e6ebb37789"
)
EXPECTED_FRAME_COUNT = 64
EXPECTED_INSTANCE_COUNT = 775
BACKGROUND_LABEL = {
    "name": "background",
    "color": "#000000",
    "attributes": [],
}


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


def _label_specs(raw_labels: list[dict[str, Any]]) -> list[dict[str, Any]]:
    labels = [
        {
            "name": item["name"],
            "color": item["color"],
            "attributes": [],
        }
        for item in raw_labels
    ]
    if not any(label["name"] == BACKGROUND_LABEL["name"] for label in labels):
        labels.insert(0, dict(BACKGROUND_LABEL))
    return labels


def _annotation_counts(task: Any) -> dict[str, int]:
    annotations = task.get_annotations()
    return {
        "shape_count": len(annotations.shapes),
        "tag_count": len(annotations.tags),
        "track_count": len(annotations.tracks),
    }


def _list_tasks_when_ready(client: Any, *, attempts: int = 60) -> list[Any]:
    for attempt in range(1, attempts + 1):
        try:
            return list(client.tasks.list())
        except ServiceException as error:
            if error.status not in {500, 502, 503} or attempt == attempts:
                raise
            time.sleep(2)
    raise AssertionError("unreachable")


def _task_record(task: Any, *, disposition: str, expected_labels: list[str]) -> dict[str, Any]:
    task.fetch()
    task_labels = list(task.get_labels())
    actual_labels = sorted(label.name for label in task_labels)
    if task.size != EXPECTED_FRAME_COUNT:
        raise RuntimeError(
            f"Task {task.id} has {task.size} frames, expected {EXPECTED_FRAME_COUNT}"
        )
    if actual_labels != sorted(expected_labels):
        raise RuntimeError(
            f"Task {task.id} labels differ: actual={actual_labels!r} "
            f"expected={sorted(expected_labels)!r}"
        )
    annotation_counts = _annotation_counts(task)
    label_names_by_id = {label.id: label.name for label in task_labels}
    annotations = task.get_annotations()
    shape_counts_by_label = dict(
        sorted(
            Counter(
                label_names_by_id.get(shape.label_id, f"unknown:{shape.label_id}")
                for shape in annotations.shapes
            ).items()
        )
    )
    return {
        "id": task.id,
        "name": task.name,
        "disposition": disposition,
        "frame_count": task.size,
        "label_names": actual_labels,
        **annotation_counts,
        "shape_counts_by_label": shape_counts_by_label,
        "url_path": f"/tasks/{task.id}",
    }


def _ensure_task(
    client: Any,
    *,
    name: str,
    labels: list[dict[str, Any]],
    images_path: Path,
    annotation_path: Path,
    annotation_format: str,
    expected_shape_count: int | None,
) -> dict[str, Any]:
    matches = [task for task in _list_tasks_when_ready(client) if task.name == name]
    if len(matches) > 1:
        raise RuntimeError(f"Multiple CVAT tasks have the reserved name {name!r}")

    expected_labels = [label["name"] for label in labels]
    if matches:
        task = matches[0]
        task.fetch()
        actual_labels = sorted(label.name for label in task.get_labels())
        expected_without_background = sorted(
            label for label in expected_labels if label != BACKGROUND_LABEL["name"]
        )
        disposition = "reused"
        if (
            actual_labels == expected_without_background
            and BACKGROUND_LABEL["name"] in expected_labels
        ):
            task.update(
                models.PatchedTaskWriteRequest(
                    labels=[models.PatchedLabelRequest(**BACKGROUND_LABEL)]
                )
            )
            disposition = "reused-and-background-label-added"

        record = _task_record(
            task,
            disposition=disposition,
            expected_labels=expected_labels,
        )
        if record["shape_count"] == 0:
            task.import_annotations(annotation_format, annotation_path)
            record = _task_record(
                task,
                disposition=f"{disposition}-and-annotations-imported",
                expected_labels=expected_labels,
            )
        if expected_shape_count is not None and record["shape_count"] != expected_shape_count:
            raise RuntimeError(
                f"Task {task.id} has {record['shape_count']} shapes, "
                f"expected {expected_shape_count}"
            )
        return record

    task = client.tasks.create_from_data(
        spec=models.TaskWriteRequest(
            name=name,
            labels=labels,
            segment_size=EXPECTED_FRAME_COUNT,
            overlap=0,
        ),
        resources=[images_path],
        resource_type=ResourceType.LOCAL,
        data_params={
            "image_quality": 100,
            "sorting_method": "lexicographical",
            "use_cache": False,
        },
        annotation_path=annotation_path,
        annotation_format=annotation_format,
        status_check_period=2,
    )
    record = _task_record(task, disposition="created", expected_labels=expected_labels)
    if expected_shape_count is not None and record["shape_count"] != expected_shape_count:
        raise RuntimeError(
            f"Task {task.id} has {record['shape_count']} shapes, "
            f"expected {expected_shape_count}"
        )
    return record


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("--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()
    if workspace.name != WORKSPACE_ID:
        raise RuntimeError(f"Unexpected annotation workspace: {workspace}")

    manifest_path = workspace / "manifest.json"
    labels_path = workspace / "cvat" / "labels.json"
    images_path = workspace / "cvat" / "images.zip"
    instance_path = workspace / "cvat" / "instance-coco.zip"
    semantic_path = workspace / "cvat" / "semantic-mask.zip"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    labels = json.loads(labels_path.read_text(encoding="utf-8"))

    if manifest["ground_truth"] is not False:
        raise RuntimeError("LAB E2 import must remain an unreviewed model draft")
    if manifest["identity"]["frame_count"] != EXPECTED_FRAME_COUNT:
        raise RuntimeError("Unexpected LAB E2 frame count")
    if manifest["identity"]["draft_instance_count"] != EXPECTED_INSTANCE_COUNT:
        raise RuntimeError("Unexpected LAB E2 instance count")

    task_specs = [
        {
            "name": "LAB E2 | K1 | instance prelabels | pack 7a983bba",
            "labels": _label_specs(labels["instance_task"]),
            "annotation_path": instance_path,
            "annotation_format": "COCO 1.0",
            "expected_shape_count": EXPECTED_INSTANCE_COUNT,
        },
        {
            "name": "LAB E2 | K1 | dense semantic prelabels | pack 7a983bba",
            "labels": _label_specs(labels["semantic_task"]),
            "annotation_path": semantic_path,
            "annotation_format": "Segmentation mask 1.1",
            "expected_shape_count": None,
        },
    ]

    with make_client(args.server, credentials=(args.username, password)) as client:
        client.check_server_version(fail_if_unsupported=True)
        task_records = [
            _ensure_task(
                client,
                name=task_spec["name"],
                labels=task_spec["labels"],
                images_path=images_path,
                annotation_path=task_spec["annotation_path"],
                annotation_format=task_spec["annotation_format"],
                expected_shape_count=task_spec["expected_shape_count"],
            )
            for task_spec in task_specs
        ]

    report = {
        "schema_version": "missioncore.lab-e2-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": WORKSPACE_ID,
        "evaluation_pack_id": EVALUATION_PACK_ID,
        "workspace_manifest_sha256": _sha256(manifest_path),
        "ground_truth": False,
        "inputs": {
            "images_zip_sha256": _sha256(images_path),
            "instance_coco_zip_sha256": _sha256(instance_path),
            "semantic_mask_zip_sha256": _sha256(semantic_path),
        },
        "tasks": task_records,
        "next_gate": (
            "two-pass human review and reviewed export; "
            "do not treat drafts as accuracy evidence"
        ),
    }
    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()
