"""Weights utils. Stuff related to weights repository, downloading and loading weights"""
import urllib.request
from pathlib import Path
from pprint import pformat
import os
import torch as tr

from ..logger import vre_logger as logger
from .utils import get_project_root, is_git_lfs, DownloadProgressBar

REPO_URL = "https://gitlab.com/video-representations-extractor/video-representations-extractor"
RESOURCES_URL = f"{REPO_URL}/-/raw/master/resources/"
RESOURCES_DIR = get_project_root() / "resources"

# TODO: this is a stupid way to do it but i can't get git-lfs info or scrap gitlab yet.
def _get_weights_files_for_representation(repr_name: str) -> list[str]:
    """This can kinda be generated by:
    `cd resources/weights && find . -type f`
    +
    ```python
    res = {}
    for item in a:
        item = item.removeprefix("./")
        repr_name, rest = "/".join(item.split("/")[0:2]), "/".join(item.split("/")[2:])
        res.setdefault(repr_name, []).append(rest)
    pprint(res)
    ```
    """
    weights_repository = {
        'depth/dpt': ['depth_dpt_midas.pth'],
        'depth/marigold': ['marigold-lcm-v1-0_unet.pt/0000.pt', 'marigold-lcm-v1-0_unet.pt/0002.pt',
                           'marigold-lcm-v1-0_unet.pt/0003.pt', 'marigold-lcm-v1-0_unet.pt/0001.pt', 'vae.pt'],
        'edges/dexined': ['dexined.pth'],
        'optical_flow/raft': ['raft-things.ckpt'],
        'optical_flow/rife': ['contextnet.ckpt', 'unet.ckpt', 'flownet.ckpt'],
        'semantic_segmentation/mask2former': ['49189528_0.ckpt', '49189528_1.ckpt', '47429163_0.ckpt'],
        'semantic_segmentation/safeuav': ['safeuav_semantic_0956_pytorch.ckpt'],
        'soft_segmentation/fastsam': ['FastSAM-x.pt', 'FastSAM-s.pt']
    }
    assert repr_name in weights_repository, f"'{repr_name}' not in {pformat(weights_repository)}"
    return weights_repository[repr_name]

def _get_weights_dir() -> Path:
    """gets the weights dir of this project"""
    return Path(os.getenv("VRE_WEIGHTS_DIR", RESOURCES_DIR / "weights"))

def fetch_resource(resource_name: str) -> Path:
    """fetches a resources from gitlab LFS if needed"""
    url = f"{RESOURCES_URL}/{resource_name}"
    path = get_project_root() / "resources" / resource_name
    if not Path(resource_name).exists() is is_git_lfs(path):
        with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=resource_name) as t:
            urllib.request.urlretrieve(url, filename=path, reporthook=t.update_to)
    return path

def fetch_weights(repr_file: str | Path) -> Path:
    """fetches weights for a representation. repr_file is expected to be Path(__file__) from the caller script"""
    repr_name = "/".join([x.name for x in Path(repr_file).parents[0:2][::-1]]) # i.e. [depth/marigold]/{marigold.py}
    files = _get_weights_files_for_representation(repr_name)
    for file in files:
        url = f"{RESOURCES_URL}/weights/{repr_name}/{file}"
        output_path = _get_weights_dir() / repr_name / file
        if output_path.exists() and not is_git_lfs(output_path):
            continue
        output_path.parent.mkdir(exist_ok=True, parents=True)
        logger.debug(f"'{repr_name}/{file}' does not exist. Downloading from remote weights repository")
        with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=f"{repr_name}/{file}") as t:
            urllib.request.urlretrieve(url, filename=output_path, reporthook=t.update_to)
    return _get_weights_dir() / repr_name

def vre_load_weights(path: Path) -> dict[str, tr.Tensor]:
    """load weights from disk. weights can be sharded as well. Sometimes we store this due to git-lfs big files bug."""
    logger.debug(f"Loading weights from '{path}'")
    if path.is_dir():
        res = {}
        for item in path.iterdir():
            res = {**res, **tr.load(item, map_location="cpu")}
        return res
    return tr.load(path, map_location="cpu")
