"""Models in the wrong models folder, and identical copies (the models page's Tidy view,
ui_model_tidy.js).

GET  /anomalous/placement/check[?deep=1] -> {misplaced, duplicates, shared, unchecked}
    misplaced: models whose header (model_kind.py, only when sure) names another type than
    the folder they are in, with where they belong (the same subfolder in a folder of that
    type, on the same drive when there is one) and whether they still load where they are.
    A file a node pack here reads from its folder (PACK_READERS) is not misplaced.
    duplicates: files with the same SHA-256 — known from scans, else read now for files of
    the same size (small ones always, large ones with deep). shared: identical files each in
    a folder of its own type where it loads (one per loader, e.g. an Anima LLLite in
    model_patches and in controlnet): both are used, none is extra. unchecked: the same-size
    large files a deep check would still read.
GET  /anomalous/placement/look?type&root&rel -> {kind, sure, base}: what one model's header says
    it is, even when unsure (the models page's "Change type", ui_model_retype.js).
POST /anomalous/placement/move {type, root, rel, to_type, to_root, to_rel} -> moves a model
    and its sidecars (covers, info, the user's notes); never over a file (" (2)"). Only for
    the computer ComfyUI runs on, unless remote changes are allowed (local_guard.may_change).

belongs() is the same check for one file (the AI's search_models and move_model,
api/mcp_model_use.py). Removing a copy goes through /anomalous/delete_model (Recycle Bin). Headers and computed
SHA-256s are kept in memory per file (path, size, time).
"""

import asyncio
import os
import threading
import urllib.parse

from aiohttp import web
import folder_paths

try:
    from ..model_identity import file_sha256
    from ..model_kind import inspect_file
    from ..model_sidecars import drop_empty_hidden_dir, existing, follow
except ImportError:  # loaded outside the package (tests)
    from model_identity import file_sha256
    from model_kind import inspect_file
    from model_sidecars import drop_empty_hidden_dir, existing, follow
from .activity_log import add_entry
from .local_guard import may_change
from .metadata import get_metadata
from .model_constants import MEDIA_EXTENSIONS, PREVIEW_SUFFIXES, SIDECAR_SUFFIXES
from .model_download import MODEL_EXTENSIONS, clean_folder, forget_file_lists
from .model_import import Refused, destination, import_types, move_into
from .path_utils import model_reference_within

# (folder type, what the file is): it loads there all the same. The checkpoint loader reads a
# file of only the drawing model too (as a diffusion model, without VAE and text encoders).
WORKS_ANYWAY = {("diffusion_models", "checkpoints"), ("checkpoints", "diffusion_models")}
# Node packs that read a kind of file from another folder than ComfyUI's own loader:
# (folder type, kind, form) -> the pack's node. Anima LLLite: ComfyUI's "Load Model Patch"
# reads model_patches, the ComfyUI-Anima-LLLite pack reads controlnet.
PACK_READERS = {("controlnet", "model_patches", "lllite"): "AnimaLLLiteApply_sdscripts"}
# Same-size files up to this are read at once (LoRAs, embeddings); larger ones on "Compare".
AUTO_HASH = 256 * 1024 * 1024
CACHE_LIMIT = 5000

_lock = threading.Lock()
_heads = {}   # path -> (size, mtime_ns, inspect_file result)
_hashes = {}  # path -> (size, mtime_ns, sha256)


def _cached(store, path, stat, compute):
    key = (stat.st_size, stat.st_mtime_ns)
    with _lock:
        hit = store.get(path)
    if hit and hit[:2] == key:
        return hit[2]
    value = compute(path)
    with _lock:
        if len(store) >= CACHE_LIMIT:
            store.clear()
        store[path] = (*key, value)
    return value


def library():
    """Every model file once: [{type, root, rel, path, size, dev}] (a folder listed under two
    types counts for the first)."""
    seen, files = set(), []
    for folder_type in import_types():
        for index, base in enumerate(folder_paths.get_folder_paths(folder_type) or []):
            if not os.path.isdir(base):
                continue
            for current, dirs, names in os.walk(base):
                dirs[:] = sorted(name for name in dirs if not name.startswith("."))
                for name in sorted(names):
                    if not name.lower().endswith(MODEL_EXTENSIONS):
                        continue
                    path = os.path.join(current, name)
                    real = os.path.realpath(path)
                    if real in seen:
                        continue
                    seen.add(real)
                    try:
                        stat = os.stat(path)
                    except OSError:
                        continue
                    files.append({"type": folder_type, "root": index, "rel": os.path.relpath(path, base).replace(os.sep, "/"),
                                  "path": path, "size": stat.st_size, "dev": stat.st_dev, "stat": stat})
    return files


def _card(item):
    """What the page shows of a file: its name (the user's, else Civitai's, else the file's),
    Civitai version and cover URL."""
    try:
        meta = get_metadata(item["path"])
    except Exception:
        meta = {}
    folder, name = os.path.split(item["path"])
    stem = os.path.splitext(name)[0]
    cover = next((stem + suffix for suffix in PREVIEW_SUFFIXES + MEDIA_EXTENSIONS
                  if os.path.isfile(os.path.join(folder, stem + suffix))), "")
    sub = item["rel"].rpartition("/")[0]
    url = ""
    if cover:
        url = (f"/anomalous/image?type={urllib.parse.quote(item['type'])}&path_idx={item['root']}"
               f"&subfolder={urllib.parse.quote(sub)}&filename={urllib.parse.quote(cover)}")
    version = meta.get("version_name") or ""
    name = meta.get("custom_name") or ""
    if not name and meta.get("model_id"):  # Civitai's "Model - Version": the model's part
        full = meta.get("name") or ""
        name = full[:-len(" - " + version)] if version and full.endswith(" - " + version) else full
    return {"name": name or stem, "version": version, "preview_url": url}


def home_for(item, kind):
    """Where a file of `kind` goes: its own subfolder and name, in a folder of that type,
    preferably on its drive (a move is then a rename)."""
    roots = []
    for index, base in enumerate(folder_paths.get_folder_paths(kind) or []):
        if os.path.isdir(base):
            roots.append((index, os.stat(base).st_dev))
    if not roots:
        return None
    index = next((index for index, dev in roots if dev == item["dev"]), roots[0][0])
    return {"type": kind, "root": index, "rel": item["rel"]}


def _node_here(name):
    import nodes  # ComfyUI's registry, filled once ComfyUI has loaded the packs
    return name in getattr(nodes, "NODE_CLASS_MAPPINGS", {})


def loads_here(folder_type, found):
    """Whether a file whose header says `found` loads from a `folder_type` folder: it is of
    that type (or the header cannot tell), or a node pack here reads it from there."""
    kind = found.get("kind")
    if not kind or kind == folder_type:
        return True
    node = PACK_READERS.get((folder_type, kind, found.get("form")))
    return bool(node) and _node_here(node)


def misplaced(files):
    found = []
    for item in files:
        result = _cached(_heads, item["path"], item["stat"], inspect_file)
        kind = result.get("kind")
        if not result.get("sure") or kind not in import_types() or loads_here(item["type"], result):
            continue
        home = home_for(item, kind)
        if home:
            found.append({"type": item["type"], "root": item["root"], "rel": item["rel"], "size": item["size"],
                          "kind": kind, "base": result.get("base", ""), "works": (item["type"], kind) in WORKS_ANYWAY,
                          "to": home, **_card(item)})
    return found


def one(folder_type, root, rel, path):
    """One model file as library() lists it."""
    stat = os.stat(path)
    return {"type": folder_type, "root": root, "rel": rel, "path": path, "size": stat.st_size, "dev": stat.st_dev, "stat": stat}


def belongs(item):
    """misplaced() of one file (from one()): where it belongs, or None when it is where it goes or its kind is unsure."""
    found = misplaced([item])
    return found[0] if found else None


def _known_hash(item):
    try:
        value = str(get_metadata(item["path"]).get("hash") or "").lower()
    except Exception:
        value = ""
    if len(value) == 64:
        return value
    with _lock:
        hit = _hashes.get(item["path"])
    return hit[2] if hit and hit[:2] == (item["stat"].st_size, item["stat"].st_mtime_ns) else ""


def _copy(item):
    """One of identical files: where it is, how it shows, whether its folder fits what it is
    (the copy to keep, when only one does)."""
    found = _cached(_heads, item["path"], item["stat"], inspect_file)
    return {"type": item["type"], "root": item["root"], "rel": item["rel"],
            "fits": loads_here(item["type"], found), **_card(item)}


def duplicates(files, deep=False):
    """(groups of identical files, groups whose every copy loads in a folder of its own type,
    {groups, bytes} still to read for a deep check)."""
    by_size = {}
    for item in files:
        if item["size"]:
            by_size.setdefault(item["size"], []).append(item)
    groups, unchecked = {}, {"groups": 0, "bytes": 0}
    for size, items in by_size.items():
        if len(items) < 2:
            continue
        unknown = [item for item in items if not _known_hash(item)]
        read = deep or size <= AUTO_HASH
        if unknown and not read:
            unchecked["groups"] += 1
            unchecked["bytes"] += size * len(unknown)
        for item in items:
            sha = _known_hash(item)
            if not sha and read:
                sha = _cached(_hashes, item["path"], item["stat"], file_sha256)
            if sha:
                groups.setdefault(sha, []).append(item)
    found, shared = [], []
    for sha, items in groups.items():
        if len(items) < 2:
            continue
        group = {"sha256": sha, "size": items[0]["size"], "files": [_copy(item) for item in items]}
        one_each = len({item["type"] for item in items}) == len(items)
        (shared if one_each and all(file["fits"] for file in group["files"]) else found).append(group)
    return found, shared, unchecked


def check(deep=False):
    files = library()
    found, shared, unchecked = duplicates(files, deep)
    return {"misplaced": misplaced(files), "duplicates": found, "shared": shared, "unchecked": unchecked, "checked": len(files)}


def _sidecars(path):
    """The model's covers, info and notes: files named by its stem, beside it or in the folder's
    .amb, unless another model in its folder has that stem too (then they are shared, and stay)."""
    folder, name = os.path.split(path)
    stem = os.path.splitext(name)[0]
    if any(os.path.isfile(os.path.join(folder, stem + ext)) for ext in MODEL_EXTENSIONS if stem + ext != name):
        return []
    return existing(os.path.join(folder, stem), SIDECAR_SUFFIXES)


def _source(folder_type, root, rel):
    """(root index, clean rel, file) of the model at type `folder_type`, root index `root`, `rel`."""
    if folder_type not in import_types():
        raise Refused("bad_type")
    paths = folder_paths.get_folder_paths(folder_type) or []
    try:
        index = int(root)
        if index < 0:
            raise IndexError(index)
        base = os.path.realpath(paths[index])
        rel = clean_folder(rel)
        source = model_reference_within(base, rel.replace("/", os.sep))
    except (ValueError, TypeError, IndexError):
        raise Refused("bad_path")
    if not os.path.isfile(source) or not rel.lower().endswith(MODEL_EXTENSIONS):
        raise Refused("gone")
    return index, rel, source


def look(folder_type, root, rel):
    """What the header of one model says it is: {type, root, rel, kind, sure, base}."""
    index, rel, source = _source(folder_type, root, rel)
    found = _cached(_heads, source, os.stat(source), inspect_file)
    return {"type": folder_type, "root": index, "rel": rel, "kind": found.get("kind") or "",
            "sure": bool(found.get("sure")), "base": found.get("base") or ""}


def move(data):
    folder_type = data.get("type")
    index, rel, source = _source(folder_type, data.get("root", 0), data.get("rel"))
    to_type, to_root, root_path, dest = destination(
        {"type": data.get("to_type"), "root": data.get("to_root", 0), "rel": data.get("to_rel")}, os.path.getsize(source))
    sidecars = _sidecars(source)
    move_into(source, dest, keep=False)
    moved, left = [os.path.basename(dest)], []
    dest_stem = os.path.splitext(dest)[0]
    for path, suffix in sidecars:
        target = follow(path, dest_stem, suffix)
        try:
            if os.path.exists(target):
                raise FileExistsError(target)
            move_into(path, target, keep=False)
            moved.append(os.path.basename(target))
        except OSError as error:
            left.append(os.path.basename(path))
            print(f"[Anomalous Browser] Tidy: {os.path.basename(path)} stayed where it was: {error}")
    drop_empty_hidden_dir(os.path.dirname(source))
    forget_file_lists()
    to_rel = os.path.relpath(dest, root_path).replace(os.sep, "/")
    add_entry("file", "model_move", to_rel, {"move": {
        "from": {"type": folder_type, "path_idx": index, "rel": rel},
        "to": {"type": to_type, "path_idx": to_root, "rel": to_rel}, "files": moved, "left": left,
    }})
    return {"type": to_type, "root": to_root, "rel": to_rel, "left": left}


async def api_check(request):
    deep = request.query.get("deep") == "1"
    return web.json_response(await asyncio.to_thread(check, deep))


async def api_look(request):
    query = request.query
    try:
        return web.json_response(await asyncio.to_thread(look, query.get("type"), query.get("root", "0"), query.get("rel")))
    except Refused as error:
        return web.json_response({"error": error.code}, status=404 if error.code == "gone" else 400)


async def api_move(request):
    if not may_change(request):
        return web.json_response({"error": "not_local"}, status=403)
    try:
        data = await request.json()
        if not isinstance(data, dict):
            raise ValueError
        return web.json_response(await asyncio.to_thread(move, data))
    except Refused as error:
        return web.json_response({"error": error.code}, status=404 if error.code == "gone" else 400)
    except ValueError:
        return web.json_response({"error": "bad_request"}, status=400)
    except OSError as error:
        print(f"[Anomalous Browser] Tidy: could not move the model: {error}")
        return web.json_response({"error": "io"}, status=500)


def register_routes(app):
    app.router.add_get("/anomalous/placement/check", api_check)
    app.router.add_get("/anomalous/placement/look", api_look)
    app.router.add_post("/anomalous/placement/move", api_move)