import logging
import torch
import os

import comfy.ops
import comfy.model_management as mm
import comfy.model_patcher
from comfy.utils import load_torch_file
import folder_paths

from .depth_anything_v3.configs import MODEL_CONFIGS
from .depth_anything_v3.model import (DepthAnything3Net, DinoV2, DualDPT, DPT,)
from .depth_anything_v3.camera import CameraEnc, CameraDec
from .depth_anything_v3.gs import GSDPT, GaussianAdapter
from .utils import DEFAULT_PATCH_SIZE, check_model_capabilities

_da3_model_dir = os.path.join(folder_paths.models_dir, "depthanything3")
os.makedirs(_da3_model_dir, exist_ok=True)
folder_paths.add_model_folder_path("depth_anything_v3", _da3_model_dir)


def _get_da3_model_list():
    return folder_paths.get_filename_list("depth_anything_v3")

ENCODER_EMBED_DIMS = {
    'vits': 384,
    'vitb': 768,
    'vitl': 1024,
    'vitg': 1536,
}

def detect_da3_variant(state_dict):
    keys = set(state_dict.keys())
    stripped_keys = set()
    for k in keys:
        stripped_keys.add(k[6:] if k.startswith('model.') else k)

    has_da3_metric = any(k.startswith('da3_metric.') for k in stripped_keys)
    if has_da3_metric:
        return 'da3nested-giant-large'

    has_da3_prefix = any(k.startswith('da3.') for k in stripped_keys)
    prefix = 'da3.' if has_da3_prefix else ''
    patch_key = f'{prefix}net.patch_embed.proj.weight'
    if patch_key in stripped_keys:
        embed_dim = state_dict.get(patch_key, state_dict.get(f'model.{patch_key}')).shape[0]
    else:
        embed_dim = None
        for k in stripped_keys:
            if 'patch_embed.proj.weight' in k:
                sd_key = k if k in state_dict else f'model.{k}'
                if sd_key in state_dict:
                    embed_dim = state_dict[sd_key].shape[0]
                break

    dim_to_encoder = {384: 'vits', 768: 'vitb', 1024: 'vitl', 1536: 'vitg'}
    encoder = dim_to_encoder.get(embed_dim)
    has_cam = any('cam_enc' in k for k in stripped_keys)
    has_gs = any('gs_head' in k for k in stripped_keys)
    head_output_key = f'{prefix}head.output_conv1.2.weight'
    is_dual_head = False
    if head_output_key in stripped_keys:
        sd_key = head_output_key if head_output_key in state_dict else f'model.{head_output_key}'
        if sd_key in state_dict:
            out_channels = state_dict[sd_key].shape[0]
            is_dual_head = (out_channels == 2)

    if encoder == 'vitg' and has_cam:
        if has_gs:
            return 'da3-giant'
        return 'da3-giant'  # giant always has GS potential
    elif encoder == 'vitl':
        if has_cam and is_dual_head:
            return 'da3-large'
        elif not has_cam and not is_dual_head:
            return 'da3mono-large'
        elif has_cam:
            return 'da3-large'
        else:
            return 'da3mono-large'
    elif encoder == 'vitb':
        return 'da3-base'
    elif encoder == 'vits':
        return 'da3-small'
    return 'da3-large'


def detect_da3_variant_with_filename_hint(state_dict, filename):
    variant = detect_da3_variant(state_dict)
    if variant == 'da3mono-large' and filename:
        fname_lower = filename.lower()
        if 'metric' in fname_lower:
            return 'da3metric-large'

    return variant

def _build_gs_modules(config, operations):
    gs_head = GSDPT(
        dim_in=config['dim_in'],
        output_dim=38,
        features=config['features'],
        out_channels=config['out_channels'],
        operations=operations,
    )

    gs_adapter = GaussianAdapter(
        sh_degree=2,
        pred_color=False,
        pred_offset_depth=True,
        pred_offset_xy=True,
        gaussian_scale_min=1e-5,
        gaussian_scale_max=30.0,
    )

    return gs_head, gs_adapter


class DA3ModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.da3 = model

    def forward(self, *args, **kwargs):
        return self.da3(*args, **kwargs)

    def to(self, *args, **kwargs):
        self.da3 = self.da3.to(*args, **kwargs)
        return self

    @property
    def cam_enc(self):
        return self.da3.cam_enc if hasattr(self.da3, 'cam_enc') else None

    @property
    def cam_dec(self):
        return self.da3.cam_dec if hasattr(self.da3, 'cam_dec') else None

    @property
    def gs_head(self):
        return self.da3.gs_head if hasattr(self.da3, 'gs_head') else None

    @property
    def gs_adapter(self):
        return self.da3.gs_adapter if hasattr(self.da3, 'gs_adapter') else None


class NestedModelWrapper(torch.nn.Module):
    def __init__(self, da3_main, da3_metric):
        super().__init__()
        self.da3 = da3_main
        self.da3_metric = da3_metric

    def forward(self, *args, **kwargs):
        from .depth_anything_v3.alignment import (
            apply_metric_scaling, compute_sky_mask, compute_alignment_mask,
            sample_tensor_for_quantile, least_squares_scale_scalar
        )

        output = self.da3(*args, **kwargs)
        x = args[0] if args else kwargs.get('x')
        infer_gs = kwargs.get('infer_gs', False)
        metric_output = self.da3_metric(x, infer_gs=infer_gs)

        metric_output.depth = apply_metric_scaling(
            metric_output.depth,
            output.intrinsics,
        )

        non_sky_mask = compute_sky_mask(metric_output.sky, threshold=0.3)

        if non_sky_mask.sum() > 10:
            depth_conf_ns = output.depth_conf[non_sky_mask]
            depth_conf_sampled = sample_tensor_for_quantile(depth_conf_ns, max_samples=100000)
            median_conf = torch.quantile(depth_conf_sampled, 0.5)

            align_mask = compute_alignment_mask(
                output.depth_conf, non_sky_mask, output.depth, metric_output.depth, median_conf
            )

            valid_depth = output.depth[align_mask]
            valid_metric_depth = metric_output.depth[align_mask]
            scale_factor = least_squares_scale_scalar(valid_metric_depth, valid_depth)

            output.depth *= scale_factor
            if hasattr(output, 'extrinsics') and output.extrinsics is not None:
                output.extrinsics[:, :, :3, 3] *= scale_factor
            output.is_metric = 1
            output.scale_factor = scale_factor.item()

            non_sky_depth = output.depth[non_sky_mask]
            if non_sky_depth.numel() > 100000:
                idx = torch.randint(0, non_sky_depth.numel(), (100000,), device=non_sky_depth.device)
                sampled_depth = non_sky_depth.flatten()[idx]
            else:
                sampled_depth = non_sky_depth
            max_depth = torch.quantile(sampled_depth, 0.99)
            sky_depth = max(200.0, max_depth.item() * 2.0)
            output.depth[~non_sky_mask] = sky_depth
            output.sky = metric_output.sky
        else:
            output.sky = metric_output.sky

        return output

    def to(self, *args, **kwargs):
        self.da3 = self.da3.to(*args, **kwargs)
        self.da3_metric = self.da3_metric.to(*args, **kwargs)
        return self

    @property
    def cam_enc(self):
        return self.da3.cam_enc if hasattr(self.da3, 'cam_enc') else None

    @property
    def cam_dec(self):
        return self.da3.cam_dec if hasattr(self.da3, 'cam_dec') else None

    @property
    def gs_head(self):
        return self.da3.gs_head if hasattr(self.da3, 'gs_head') else None

    @property
    def gs_adapter(self):
        return self.da3.gs_adapter if hasattr(self.da3, 'gs_adapter') else None


def _build_da3_model(model_path, model_key, dtype, attention):
    config = MODEL_CONFIGS[model_key]
    is_nested = config.get('is_nested', False)
    operations = comfy.ops.manual_cast

    with torch.device("meta"):
        if is_nested:
            backbone_main = DinoV2(
                name=config['encoder'],
                out_layers=config.get('out_layers', [19, 27, 33, 39]),
                alt_start=config.get('alt_start', 13),
                qknorm_start=config.get('qknorm_start', 13),
                rope_start=config.get('rope_start', 13),
                cat_token=config.get('cat_token', True),
                operations=operations,
            )
            head_main = DualDPT(
                dim_in=config['dim_in'],
                output_dim=2,
                features=config['features'],
                out_channels=config['out_channels'],
                operations=operations,
            )
            embed_dim = ENCODER_EMBED_DIMS.get(config['encoder'], 1536)
            cam_enc_main = CameraEnc(
                dim_out=embed_dim,
                dim_in=9,
                trunk_depth=4,
                num_heads=embed_dim // 64,
                mlp_ratio=4,
                init_values=0.01,
                operations=operations,
            )
            cam_dec_main = CameraDec(dim_in=config['dim_in'], operations=operations)
            gs_head_main, gs_adapter_main = _build_gs_modules(config, operations)

            da3_main = DepthAnything3Net(
                net=backbone_main,
                head=head_main,
                cam_dec=cam_dec_main,
                cam_enc=cam_enc_main,
                gs_head=gs_head_main,
                gs_adapter=gs_adapter_main,
            )

            metric_config = MODEL_CONFIGS.get('da3metric-large', {
                'encoder': 'vitl',
                'features': 256,
                'out_channels': [256, 512, 1024, 1024],
                'dim_in': 1024,
                'out_layers': [4, 11, 17, 23],
            })
            backbone_metric = DinoV2(
                name=metric_config.get('encoder', 'vitl'),
                out_layers=metric_config.get('out_layers', [4, 11, 17, 23]),
                alt_start=-1,
                qknorm_start=-1,
                rope_start=-1,
                cat_token=False,
                operations=operations,
            )
            head_metric = DPT(
                dim_in=metric_config.get('dim_in', 1024),
                output_dim=1,
                features=metric_config.get('features', 256),
                out_channels=metric_config.get('out_channels', [256, 512, 1024, 1024]),
                operations=operations,
            )
            da3_metric = DepthAnything3Net(
                net=backbone_metric,
                head=head_metric,
                cam_dec=None,
                cam_enc=None,
                gs_head=None,
                gs_adapter=None,
            )

            inner_model = NestedModelWrapper(da3_main, da3_metric)
        else:
            backbone = DinoV2(
                name=config['encoder'],
                out_layers=config.get('out_layers', [4, 11, 17, 23]),
                alt_start=config.get('alt_start', -1),
                qknorm_start=config.get('qknorm_start', -1),
                rope_start=config.get('rope_start', -1),
                cat_token=config.get('cat_token', False),
                operations=operations,
            )

            if config.get('is_mono', False) or config.get('is_metric', False):
                head = DPT(
                    dim_in=config['dim_in'],
                    output_dim=1,
                    features=config['features'],
                    out_channels=config['out_channels'],
                    operations=operations,
                )
            else:
                head = DualDPT(
                    dim_in=config['dim_in'],
                    output_dim=2,
                    features=config['features'],
                    out_channels=config['out_channels'],
                    operations=operations,
                )

            cam_enc = None
            cam_dec = None
            if config.get('has_cam', False) and config.get('alt_start', -1) != -1:
                embed_dim = ENCODER_EMBED_DIMS.get(config['encoder'], 1024)
                cam_enc = CameraEnc(
                    dim_out=embed_dim,
                    dim_in=9,
                    trunk_depth=4,
                    num_heads=embed_dim // 64,
                    mlp_ratio=4,
                    init_values=0.01,
                    operations=operations,
                )
                cam_dec = CameraDec(dim_in=config['dim_in'], operations=operations)

            gs_head = None
            gs_adapter = None
            if config.get('has_3d_gaussians', model_key == 'da3-giant'):
                gs_head, gs_adapter = _build_gs_modules(config, operations)

            inner_model = DepthAnything3Net(
                net=backbone,
                head=head,
                cam_dec=cam_dec,
                cam_enc=cam_enc,
                gs_head=gs_head,
                gs_adapter=gs_adapter,
            )

    state_dict = load_torch_file(model_path)
    new_state_dict = {}
    stripped_count = 0
    for key, value in state_dict.items():
        new_key = key
        if new_key.startswith('model.'):
            new_key = new_key[6:]
            stripped_count += 1
        new_state_dict[new_key] = value

    has_da3_prefix = any(k.startswith('da3.') for k in new_state_dict.keys())

    if is_nested:
        model = inner_model
    elif has_da3_prefix:
        model = DA3ModelWrapper(inner_model)
    else:
        model = inner_model

    expanded = {}
    for key in list(new_state_dict.keys()):
        if 'output_conv2_aux.0.' in key:
            prefix, suffix = key.split('output_conv2_aux.0.')
            for idx in range(1, 4):
                clone_key = f"{prefix}output_conv2_aux.{idx}.{suffix}"
                if clone_key not in new_state_dict:
                    expanded[clone_key] = new_state_dict[key].clone()
    if expanded:
        new_state_dict.update(expanded)

    result = model.load_state_dict(new_state_dict, strict=False, assign=True)
    model.to(dtype=dtype)

    for name, param in list(model.named_parameters()):
        if param.device.type == 'meta':
            parts = name.split('.')
            parent = model
            for part in parts[:-1]:
                parent = getattr(parent, part)
            setattr(parent, parts[-1], torch.nn.Parameter(
                torch.zeros(param.shape, dtype=dtype), requires_grad=False
            ))

    for name, buf in list(model.named_buffers()):
        if buf.device.type == 'meta':
            parts = name.split('.')
            parent = model
            for part in parts[:-1]:
                parent = getattr(parent, part)
            parent.register_buffer(parts[-1], torch.zeros(buf.shape, dtype=buf.dtype))

    model.eval()
    return model


class DownloadAndLoadDepthAnythingV3Model():
    @classmethod
    def execute(cls, model, precision="auto", attention="auto"):
        device = mm.get_torch_device()
        if precision == "auto":
            if mm.should_use_bf16(device):
                dtype = torch.bfloat16
            elif mm.should_use_fp16(device):
                dtype = torch.float16
            else:
                dtype = torch.float32
        elif precision == "bf16":
            dtype = torch.bfloat16
        elif precision == "fp16":
            dtype = torch.float16
        else:
            dtype = torch.float32

        model_path = folder_paths.get_full_path("depth_anything_v3", model)

        if model_path is None:
            raise FileNotFoundError(
                f"Model '{model}' not found in [ComfyUI]/models/depthanything3/."
            )

        sd = load_torch_file(model_path)
        model_key = detect_da3_variant_with_filename_hint(sd, model)
        del sd  # Free memory before building model
        config = MODEL_CONFIGS[model_key]
        loaded_model = _build_da3_model(model_path, model_key, dtype, attention)

        patcher = comfy.model_patcher.ModelPatcher(
            loaded_model,
            load_device=mm.get_torch_device(),
            offload_device=mm.unet_offload_device(),
        )

        patcher.model_options["da3_capabilities"] = check_model_capabilities(loaded_model)
        patcher.model_options["da3_config"] = config
        patcher.model_options["da3_dtype"] = dtype
        patcher.model_options["da3_model_key"] = model_key

        return patcher