"""Asset storage: files land in data/assets/<id>/<name>, kinds inferred by extension."""
import json
import shutil
from pathlib import Path

from . import db

ASSETS_DIR = db.DATA / "assets"

KIND_BY_EXT = {
    "video": {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v"},
    "image": {".png", ".jpg", ".jpeg", ".webp", ".heic", ".bmp", ".tif", ".tiff"},
    "model": {".glb", ".gltf", ".obj", ".fbx", ".blend", ".usd", ".usdz", ".usda",
              ".ply", ".stl", ".spz", ".abc", ".dae", ".bvh"},
}


def kind_of(name: str) -> str:
    ext = Path(name).suffix.lower()
    for kind, exts in KIND_BY_EXT.items():
        if ext in exts:
            return kind
    return "other"


def register_file(con, src: Path, name: str | None = None, parent_job: str | None = None,
                  move: bool = False, meta: dict | None = None) -> dict:
    """Copy/move a file (or directory) into the asset store and record it."""
    name = name or src.name
    asset_id = db.new_id()
    dest_dir = ASSETS_DIR / asset_id
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest = dest_dir / name
    meta = dict(meta or {})
    if src.is_dir():
        if move:
            shutil.move(str(src), dest)
        else:
            shutil.copytree(src, dest)
        size = sum(f.stat().st_size for f in dest.rglob("*") if f.is_file())
        kind = "frames" if any(dest.glob("*.jpg")) or any(dest.glob("*.png")) else "folder"
    else:
        if move:
            shutil.move(str(src), dest)
        else:
            shutil.copy2(src, dest)
        size = dest.stat().st_size
        kind = kind_of(name)
    # operators may tag a semantic kind that isn't inferable from the extension
    # (e.g. a splat .ply vs a mesh .ply, or a colmap_dataset folder)
    kind = meta.pop("kind", kind)
    con.execute(
        "INSERT INTO assets (id, name, kind, path, size, meta, parent_job, created_at) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
        (asset_id, name, kind, str(dest), size, json.dumps(meta or {}), parent_job, db.now()),
    )
    con.commit()
    return get_asset(con, asset_id)


def register_upload(con, name: str, data: bytes) -> dict:
    asset_id = db.new_id()
    dest_dir = ASSETS_DIR / asset_id
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest = dest_dir / name
    dest.write_bytes(data)
    con.execute(
        "INSERT INTO assets (id, name, kind, path, size, meta, parent_job, created_at) "
        "VALUES (?, ?, ?, ?, ?, '{}', NULL, ?)",
        (asset_id, name, kind_of(name), str(dest), len(data), db.now()),
    )
    con.commit()
    return get_asset(con, asset_id)


def get_asset(con, asset_id: str) -> dict | None:
    row = con.execute("SELECT * FROM assets WHERE id = ?", (asset_id,)).fetchone()
    return db.row_to_dict(row) if row else None


def list_assets(con) -> list[dict]:
    rows = con.execute("SELECT * FROM assets ORDER BY created_at DESC").fetchall()
    return [db.row_to_dict(r) for r in rows]


def delete_asset(con, asset_id: str) -> bool:
    asset = get_asset(con, asset_id)
    if not asset:
        return False
    shutil.rmtree(ASSETS_DIR / asset_id, ignore_errors=True)
    con.execute("DELETE FROM assets WHERE id = ?", (asset_id,))
    con.commit()
    return True