from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from .config import settings
from .paths import rel


SEGMODEL_ARCHS = [
    "Unet",
    "UnetPlusPlus",
    "FPN",
    "PSPNet",
    "DeepLabV3",
    "DeepLabV3Plus",
    "Linknet",
    "MAnet",
    "PAN",
    "UPerNet",
    "Segformer",
    "DPT",
]


YOLO_MODELS = [
    "YOLOv8n-seg",
    "YOLOv8s-seg",
    "YOLOv8m-seg",
    "YOLOv8l-seg",
    "YOLOv8x-seg",
    "YOLOv9c-seg",
    "YOLOv9e-seg",
    "YOLO11n-seg",
    "YOLO11s-seg",
    "YOLO11m-seg",
    "YOLO11l-seg",
    "YOLO11x-seg",
    "YOLO12-seg",
]


TASK_TYPES = [
    "mock.echo",
    "system.backup",
    "dataset.rename",
    "dataset.to_png",
    "dataset.resize",
    "dataset.pair",
    "dataset.rebuild_labels",
    "dataset.stack",
    "dataset.stitch",
    "dataset.video_frames",
    "segmodel.train",
    "segmodel.batch_train",
    "segmodel.predict",
    "segmodel.batch_predict",
    "segmodel.flops",
    "segmodel.raw_mask_check",
    "segmodel.metrics",
    "yolo.train",
    "yolo.batch_train",
    "yolo.predict",
    "yolo.batch_predict",
    "yolo.heatmap",
    "yolo.compare",
    "yolo.raw_mask_check",
    "yolo.video_visible",
    "yolo.video_unvisible",
    "mmseg.init_weights",
    "mmseg.generate_data",
    "mmseg.generate_alg",
    "mmseg.train",
    "mmseg.metrics",
    "mmseg.flops_fps",
    "mmseg.draw",
    "mmseg.extract_loss_miou",
    "analysis.all",
]


def _read_json(path: Path) -> Any | None:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None


def discover_datasets() -> list[dict[str, Any]]:
    root = settings.source_root
    candidates: list[dict[str, Any]] = []
    for base in ["DataSet_Public", "BestMode_Predict_Results_DataSet_Public", "Hardisk"]:
        parent = root / base
        if not parent.exists():
            continue
        for item in sorted(parent.iterdir()):
            if item.is_dir():
                candidates.append({"name": item.name, "path": rel(item, root), "source": base})
    mmseg_params = root / "Seg_All_In_One_MMSeg" / "My_All_In_One" / "1_Data_Parameter"
    for item in sorted(mmseg_params.glob("*.json")):
        data = _read_json(item)
        if item.name == "All_Data_Record.json" or not data:
            continue
        candidates.append({"name": item.stem, "path": rel(item, root), "source": "mmseg_parameter"})
    uploaded_root = settings.project_root / "var" / "uploads" / "datasets"
    if uploaded_root.exists():
        for item in sorted(uploaded_root.iterdir()):
            if item.is_dir():
                candidates.append({"name": item.name, "path": rel(item, settings.project_root), "source": "uploaded"})
    return candidates


def discover_mmseg_algorithms() -> list[str]:
    alg_dir = settings.source_root / "Seg_All_In_One_MMSeg" / "My_All_In_One" / "2_Alg_Program"
    if not alg_dir.exists():
        return []
    return sorted(path.stem for path in alg_dir.glob("*.py"))


def discover_weights_summary() -> dict[str, Any]:
    manifest = settings.weights_root / "manifest.json"
    if not manifest.exists():
        return {"manifest": None, "count": 0, "total_bytes": 0}
    data = _read_json(manifest) or {}
    return {
        "manifest": rel(manifest, settings.project_root),
        "count": len(data.get("files", [])),
        "total_bytes": data.get("total_bytes", 0),
        "updated_at": data.get("updated_at"),
    }


def get_catalog() -> dict[str, Any]:
    return {
        "source_root": str(settings.source_root),
        "project_root": str(settings.project_root),
        "task_types": TASK_TYPES,
        "segmodel_architectures": SEGMODEL_ARCHS,
        "yolo_models": YOLO_MODELS,
        "mmseg_algorithms": discover_mmseg_algorithms(),
        "datasets": discover_datasets(),
        "weights": discover_weights_summary(),
    }