import functools
import platform
import re
import subprocess
import sys

import bpy

from .runtime import runtime
from . import model_catalog as shared_model_catalog


BEN2_MODEL_KEY = "BEN2_BASE"
DEFAULT_DA3_MODEL_KEY = "BASE"
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
REQUIRED_MODEL_KEYS = (DEFAULT_MOGE2_MODEL_KEY, BEN2_MODEL_KEY)


DEVICE_LABELS = {
    "cuda": ("CUDA", "NVIDIA CUDA"),
    "directml": ("DirectML", "Windows GPU (DirectML)"),
    "mps": ("MPS", "Apple Metal (MPS)"),
    "cpu": ("CPU", "CPU"),
}


APPLE_SILICON_PATTERN = re.compile(r"\bApple M(\d+)\b", re.IGNORECASE)


def parse_apple_silicon_generation(brand):
    match = APPLE_SILICON_PATTERN.search(str(brand))
    return int(match.group(1)) if match else None


@functools.lru_cache(maxsize=1)
def apple_silicon_generation():
    if sys.platform != "darwin" or platform.machine().lower() not in {
        "arm64",
        "aarch64",
    }:
        return None
    try:
        result = subprocess.run(
            ("sysctl", "-n", "machdep.cpu.brand_string"),
            check=True,
            capture_output=True,
            text=True,
            timeout=2,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    return parse_apple_silicon_generation(result.stdout)


def supports_mps(generation=None):
    if generation is None:
        generation = apple_silicon_generation()
    return generation is not None and generation >= 4


def supported_devices():
    devices = runtime.environment_manifest().get("devices")
    if isinstance(devices, list) and devices:
        names = [str(device).lower() for device in devices]
        if not supports_mps():
            names = [name for name in names if name != "mps"]
        return names or ["cpu"]
    if sys.platform == "win32":
        return ["directml", "cpu"]
    return ["mps", "cpu"] if supports_mps() else ["cpu"]


def device_enum_items(_owner, _context):
    return [
        (name.upper(), *DEVICE_LABELS[name], index)
        for index, name in enumerate(supported_devices())
    ]


def models_status(refresh=False):
    status = runtime.server_status().get("resources", {}).get("model_manager", {})
    return {
        "catalog": tuple(
            shared_model_catalog.model_record(
                model,
                _model_files_ready(model),
            )
            for model in shared_model_catalog.DOWNLOADABLE_MODELS.values()
        ),
        "loaded": status.get("loaded", {}),
    }


def _model_files_ready(model):
    directory = runtime.storage_root() / "models" / model.directory_name
    return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)


def model_ready(model_key):
    return any(
        model.get("key") == model_key and model.get("ready")
        for model in models_status().get("catalog", ())
    )


def required_models_ready():
    return all(model_ready(model_key) for model_key in REQUIRED_MODEL_KEYS)


def model_catalog():
    return models_status().get("catalog", ())


def model_record(model_key):
    return next(
        (model for model in model_catalog() if model.get("key") == model_key),
        None,
    )


def _model_items(models):
    return [
        (
            model["key"],
            model["label"] if model.get("ready") else f"{model['label']} (Missing)",
            f"{model['description']}; license: {model['license']}",
            index,
        )
        for index, model in enumerate(models)
    ]


def model_enum_items(_owner, _context):
    return _model_items(
        model
        for model in model_catalog()
        if model.get("family") == "da3" and model.get("enabled")
    )


def moge2_model_enum_items(_owner, _context):
    return _model_items(
        model for model in model_catalog() if model.get("family") == "moge2"
    )


def balloon_depth_model_enum_items(_owner, _context):
    models = [
        model
        for model in model_catalog()
        if model.get("family") == "moge2" and model.get("enabled")
    ]
    models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
    return _model_items(models)

def depth_plane_model_enum_items(_owner, _context):
    models = [
        model
        for model in model_catalog()
        if model.get("family") in {"moge2", "da3"} and model.get("enabled")
    ]
    models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
    items = _model_items(models)
    items.append(
        ("ORIGINAL", "Original", "Use the original image as depth", len(items))
    )
    return items


def surface_model_enum_items(_owner, _context):
    models = [
        model
        for model in model_catalog()
        if model.get("family") in {"moge2", "da3"} and model.get("enabled")
    ]
    models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
    return _model_items(models)


BALLOON_MODE_ITEMS = (
    (
        "POISSON",
        "Poisson",
        "Use a global Poisson field for smooth inflation",
    ),
    (
        "TEDDY",
        "Teddy / Chordal Axis",
        "Use a Teddy-style chordal axis and local oval sections",
    ),
)


def balloon_mode_property():
    return bpy.props.EnumProperty(
        name="Mode",
        description="Algorithm used to inflate the saved lasso region",
        items=BALLOON_MODE_ITEMS,
        default="POISSON",
    )


def balloon_thickness_property():
    return bpy.props.FloatProperty(
        name="Thickness",
        description="Scale the Balloon thickness while keeping the lasso region",
        default=1.0,
        min=0.05,
        soft_max=5.0,
    )


def balloon_resolution_property():
    return bpy.props.FloatProperty(
        name="Resolution (m)",
        description=(
            "Distance in meters between sampled mesh points; smaller values "
            "create more geometry"
        ),
        default=0.1,
        min=0.005,
        soft_min=0.02,
        soft_max=0.5,
        max=10.0,
    )


def balloon_reuse_material_property():
    return bpy.props.BoolProperty(
        name="Reuse Material",
        description="Reuse a compatible Balloon material made from the same image",
        default=True,
    )


def balloon_contour_analysis_property():
    return bpy.props.BoolProperty(
        name="Contour",
        description=(
            "Use BEN2 Alpha above 0.5 to analyze the lasso contour before "
            "building the Balloon"
        ),
        default=False,
    )


def balloon_depth_analysis_property():
    return bpy.props.BoolProperty(
        name="Depth",
        description=("Use MoGe-2 to add adjustable image depth to the Balloon mesh"),
        default=False,
    )


def balloon_depth_model_property():
    return bpy.props.EnumProperty(
        name="Depth Model",
        description="MoGe-2 model used for Balloon depth",
        items=balloon_depth_model_enum_items,
        default=0,
    )

def balloon_radial_zoom_property(scale):
    return bpy.props.FloatVectorProperty(
        size=2,
        default=(scale, scale),
        min=0.01,
        options={"HIDDEN", "SKIP_SAVE"},
    )


class SnapFormSettings(bpy.types.PropertyGroup):
    balloon_mode: balloon_mode_property()
    balloon_thickness: balloon_thickness_property()
    balloon_resolution: balloon_resolution_property()
    balloon_reuse_material: balloon_reuse_material_property()
    balloon_contour_analysis: balloon_contour_analysis_property()
    balloon_depth_analysis: balloon_depth_analysis_property()
    balloon_depth_model: balloon_depth_model_property()
    balloon_thickness_radial_zoom: balloon_radial_zoom_property(100.0)
    balloon_resolution_radial_zoom: balloon_radial_zoom_property(1000.0)
    input_path: bpy.props.StringProperty(
        name="Image or Video",
        description="A single image or video file used for reconstruction",
        subtype="FILE_PATH",
    )
    background_input_path: bpy.props.StringProperty(
        name="Image",
        description="A single image used to test BEN2 background removal",
        subtype="FILE_PATH",
    )
    video_frames: bpy.props.IntProperty(
        name="Extract Frames",
        description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
        default=16,
        min=1,
        soft_max=64,
    )
    model: bpy.props.EnumProperty(
        name="Model",
        description="Depth Anything 3 model",
        items=model_enum_items,
    )
    device: bpy.props.EnumProperty(
        name="Debug Device",
        description=(
            "Inference device used by Debug tasks only; production tasks use "
            "the Device setting in the Add-on Preferences"
        ),
        items=device_enum_items,
        default=0,
    )
    process_res: bpy.props.IntProperty(
        name="Process Res",
        description="DA3 base processing resolution",
        default=504,
        min=1,
        soft_max=2048,
    )
    process_res_method: bpy.props.EnumProperty(
        name="Process Res Method",
        description="DA3 image resize method",
        items=(
            (
                "upper_bound_resize",
                "Upper Bound Resize",
                "Make Process Res the longer side",
            ),
            (
                "lower_bound_resize",
                "Lower Bound Resize",
                "Make Process Res the shorter side",
            ),
        ),
        default="upper_bound_resize",
    )
    moge_resolution_level: bpy.props.IntProperty(
        name="Resolution Level",
        description="MoGe-2 inference resolution level; higher is finer and slower",
        default=5,
        min=0,
        max=9,
    )
    moge_model: bpy.props.EnumProperty(
        name="Model",
        description="MoGe-2 model used for depth estimation",
        items=moge2_model_enum_items,
        default=0,
    )
    ref_view_strategy: bpy.props.EnumProperty(
        name="Ref View Strategy",
        description="DA3 reference view selection strategy",
        items=(
            ("first", "First", "Always use the first view"),
            ("middle", "Middle", "Use the middle view"),
            (
                "saddle_balanced",
                "Saddle Balanced",
                "Select a view with balanced feature metrics",
            ),
            (
                "saddle_sim_range",
                "Saddle Sim Range",
                "Select the view with the largest similarity range",
            ),
        ),
        default="saddle_balanced",
    )