"""
WorkflowUI plugin: ComfyUI API routes for media access and deletion.

- POST /delete and POST /workflowui/media/delete — delete a file by filename/subfolder/type
  (WorkflowUI calls POST /delete with JSON body).
- GET /workflowui/media/view — serve a file from output/input/temp; ?preview=webp|jpeg|SIZE for compressed image previews.
- GET /workflowui/media/list — list files in output/input/temp (optional subfolder).
- GET /workflowui/media/tree — tree of subfolders and files for browsing.
- GET /workflowui/media/capabilities — report supported features (delete, list, view).
- GET /workflowui/version_info — installed custom node modules and versions (for run metadata).
"""

from __future__ import annotations

import io
import logging
import mimetypes
import os
from pathlib import Path

from aiohttp import web

logger = logging.getLogger(__name__)

MEDIA_TYPES = ("output", "input", "temp")

PREVIEW_MAX_DIM_DEFAULT = 512
PREVIEW_QUALITY_DEFAULT = 85
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}


def _get_media_root(folder_type: str):
    """Resolve ComfyUI media root for output/input/temp. Returns None if invalid."""
    try:
        import folder_paths
        out = folder_paths.get_directory_by_type(folder_type)
        return os.path.abspath(out) if out else None
    except Exception:
        return None


def _resolve_media_path(folder_type: str, subfolder: str, filename: str) -> tuple[str | None, str | None]:
    """
    Resolve full_path for a media file. Ensures path is under base (no traversal).
    Returns (full_path, None) on success; (None, error_message) on validation failure.
    """
    if folder_type not in MEDIA_TYPES:
        return None, "Invalid type; use output, input, or temp"
    base = _get_media_root(folder_type)
    if not base:
        return None, "Could not resolve media directory"
    # Normalize subfolder: no leading/trailing slashes, no '..'
    sub = (subfolder or "").strip().replace("\\", "/").strip("/")
    if ".." in sub or sub.startswith("/"):
        return None, "Invalid subfolder"
    # Filename must be a single segment (no path)
    name = (filename or "").strip()
    if not name or "/" in name or "\\" in name or ".." in name:
        return None, "Invalid filename"
    full = os.path.join(base, sub, name) if sub else os.path.join(base, name)
    full = os.path.normpath(os.path.abspath(full))
    if not full.startswith(base):
        return None, "Path outside media directory"
    return full, None


def _delete_file(folder_type: str, subfolder: str, filename: str) -> tuple[bool, str, int]:
    """Delete one file. Returns (success, message, http_status)."""
    full, err = _resolve_media_path(folder_type, subfolder, filename)
    if err:
        return False, err, 400
    if not os.path.isfile(full):
        return False, "File not found", 404
    try:
        os.remove(full)
        return True, "ok", 200
    except OSError as e:
        logger.warning("Delete failed for %s: %s", full, e)
        return False, str(e), 500


def _is_image_file(filename: str, content_type: str | None) -> bool:
    if content_type and content_type.startswith("image/"):
        return True
    ext = Path(filename).suffix.lower()
    return ext in IMAGE_EXTENSIONS


def _make_image_preview(
    full_path: str,
    fmt: str,
    max_dim: int,
    quality: int,
) -> tuple[bytes | None, str | None]:
    """Return (body, content_type) for a resized/compressed preview, or (None, None) on failure."""
    try:
        from PIL import Image
    except ImportError:
        return None, None
    try:
        resample = getattr(Image.Resampling, "LANCZOS", getattr(Image, "LANCZOS", Image.BICUBIC))
        with Image.open(full_path) as img:
            img.load()
            if img.mode in ("RGBA", "P"):
                img = img.convert("RGB")
            w, h = img.size
            if w > max_dim or h > max_dim:
                if w >= h:
                    new_w, new_h = max_dim, max(1, int(h * max_dim / w))
                else:
                    new_w, new_h = max(1, int(w * max_dim / h)), max_dim
                img = img.resize((new_w, new_h), resample)
            buf = io.BytesIO()
            if fmt == "webp":
                img.save(buf, "WEBP", quality=quality, method=6)
                return buf.getvalue(), "image/webp"
            if fmt == "jpeg" or fmt == "jpg":
                img.save(buf, "JPEG", quality=quality, optimize=True)
                return buf.getvalue(), "image/jpeg"
            return None, None
    except Exception as e:
        logger.debug("Preview generation failed for %s: %s", full_path, e)
        return None, None


async def _handle_view(request: web.Request) -> web.Response:
    """GET ?filename=...&subfolder=...&type=output&preview=webp|jpeg|SIZE — serve file; optional compressed image preview."""
    filename = request.query.get("filename", "").strip()
    subfolder = request.query.get("subfolder", "")
    folder_type = (request.query.get("type") or "output").strip().lower()
    preview = request.query.get("preview", "").strip().lower()
    if not filename:
        return web.json_response({"error": "Missing filename"}, status=400)
    if folder_type not in MEDIA_TYPES:
        return web.json_response({"error": "Invalid type"}, status=400)
    full, err = _resolve_media_path(folder_type, subfolder, filename)
    if err:
        return web.json_response({"error": err}, status=400)
    if not os.path.isfile(full):
        return web.Response(status=404)
    content_type, _ = mimetypes.guess_type(filename, strict=False)
    if not content_type:
        content_type = "application/octet-stream"

    # Optional preview: only for images; format webp (default) or jpeg, optional max dimension
    if preview and _is_image_file(filename, content_type):
        max_dim = PREVIEW_MAX_DIM_DEFAULT
        quality = PREVIEW_QUALITY_DEFAULT
        fmt = "webp"
        if preview in ("webp", "jpeg", "jpg"):
            fmt = "jpeg" if preview in ("jpeg", "jpg") else "webp"
        else:
            try:
                max_dim = min(2048, max(64, int(preview)))
            except ValueError:
                pass
        body, preview_content_type = _make_image_preview(full, fmt, max_dim, quality)
        if body is not None and preview_content_type:
            response = web.Response(body=body)
            response.headers["Content-Type"] = preview_content_type
            response.headers["Cache-Control"] = "public, max-age=3600"
            return response

    response = web.FileResponse(full)
    response.headers["Content-Type"] = content_type
    response.headers["Cache-Control"] = "public, max-age=3600"
    return response


async def _handle_delete(request: web.Request) -> web.Response:
    """POST body: JSON { filename, subfolder?, type? }. type defaults to 'output'."""
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)
    filename = body.get("filename")
    if not filename:
        return web.json_response({"error": "Missing filename"}, status=400)
    subfolder = body.get("subfolder", "")
    folder_type = (body.get("type") or "output").strip().lower()
    if folder_type not in MEDIA_TYPES:
        return web.json_response({"error": "Invalid type"}, status=400)
    ok, msg, status = _delete_file(folder_type, subfolder, filename)
    if ok:
        return web.json_response({"ok": True}, status=status)
    return web.json_response({"ok": False, "error": msg}, status=status)


def _list_dir(folder_type: str, subfolder: str) -> tuple[list | None, str | None]:
    """List files in type/subfolder. Returns (list of {name, size, mtime}, error)."""
    if folder_type not in MEDIA_TYPES:
        return None, "Invalid type"
    base = _get_media_root(folder_type)
    if not base:
        return None, "Could not resolve media directory"
    sub = (subfolder or "").strip().replace("\\", "/").strip("/")
    if ".." in sub or sub.startswith("/"):
        return None, "Invalid subfolder"
    dir_path = os.path.join(base, sub) if sub else base
    dir_path = os.path.normpath(os.path.abspath(dir_path))
    if not dir_path.startswith(base):
        return None, "Path outside media directory"
    if not os.path.isdir(dir_path):
        return [], None
    result = []
    try:
        for name in os.listdir(dir_path):
            full = os.path.join(dir_path, name)
            if os.path.isfile(full):
                try:
                    stat = os.stat(full)
                    result.append({
                        "filename": name,
                        "size": stat.st_size,
                        "mtime": int(stat.st_mtime),
                    })
                except OSError:
                    continue
    except OSError as e:
        return None, str(e)
    result.sort(key=lambda x: (-x["mtime"], x["filename"]))
    return result, None


async def _handle_list(request: web.Request) -> web.Response:
    """GET ?type=output&subfolder= (subfolder optional)."""
    folder_type = (request.query.get("type") or "output").strip().lower()
    subfolder = request.query.get("subfolder", "")
    data, err = _list_dir(folder_type, subfolder)
    if err:
        return web.json_response({"error": err}, status=400 if "Invalid" in err else 500)
    return web.json_response({"type": folder_type, "subfolder": subfolder, "files": data})


def _tree_dir(folder_type: str, subfolder: str) -> tuple[dict | None, str | None]:
    """Return tree { folders: [name], files: [{filename, size, mtime}] } for one level."""
    if folder_type not in MEDIA_TYPES:
        return None, "Invalid type"
    base = _get_media_root(folder_type)
    if not base:
        return None, "Could not resolve media directory"
    sub = (subfolder or "").strip().replace("\\", "/").strip("/")
    if ".." in sub or sub.startswith("/"):
        return None, "Invalid subfolder"
    dir_path = os.path.join(base, sub) if sub else base
    dir_path = os.path.normpath(os.path.abspath(dir_path))
    if not dir_path.startswith(base):
        return None, "Path outside media directory"
    if not os.path.isdir(dir_path):
        return {"folders": [], "files": []}, None
    folders = []
    files = []
    try:
        for name in sorted(os.listdir(dir_path)):
            full = os.path.join(dir_path, name)
            if os.path.isdir(full):
                folders.append(name)
            elif os.path.isfile(full):
                try:
                    stat = os.stat(full)
                    files.append({"filename": name, "size": stat.st_size, "mtime": int(stat.st_mtime)})
                except OSError:
                    pass
    except OSError as e:
        return None, str(e)
    files.sort(key=lambda x: (-x["mtime"], x["filename"]))
    return {"folders": folders, "files": files}, None


async def _handle_tree(request: web.Request) -> web.Response:
    """GET ?type=output&subfolder= (one level: folders + files)."""
    folder_type = (request.query.get("type") or "output").strip().lower()
    subfolder = request.query.get("subfolder", "")
    data, err = _tree_dir(folder_type, subfolder)
    if err:
        return web.json_response({"error": err}, status=400 if "Invalid" in err else 500)
    return web.json_response({
        "type": folder_type,
        "subfolder": subfolder,
        "folders": data["folders"],
        "files": data["files"],
    })


async def _handle_capabilities(request: web.Request) -> web.Response:
    """GET /workflowui/media/capabilities — report what this plugin supports."""
    return web.json_response({
        "workflowui_plugin": True,
        "delete": True,
        "list": True,
        "tree": True,
        "view": True,
        "view_preview": True,
    })


async def _handle_version_info(request: web.Request) -> web.Response:
    """GET /workflowui/version_info — installed custom node modules and versions for run metadata."""
    from .version_info import get_version_info_payload
    return web.json_response(get_version_info_payload())


def register_routes():
    """Register routes with ComfyUI PromptServer. Call from __init__.py."""
    try:
        from server import PromptServer
        routes = PromptServer.instance.routes
    except Exception as e:
        logger.warning("WorkflowUIPlugin: could not get PromptServer routes: %s", e)
        return
    # WorkflowUI backend calls POST /delete with JSON { filename, subfolder, type }
    routes.post("/delete")(_handle_delete)
    routes.post("/workflowui/media/delete")(_handle_delete)
    routes.get("/workflowui/media/view")(_handle_view)
    routes.get("/workflowui/media/list")(_handle_list)
    routes.get("/workflowui/media/tree")(_handle_tree)
    routes.get("/workflowui/media/capabilities")(_handle_capabilities)
    routes.get("/workflowui/version_info")(_handle_version_info)
    logger.info("WorkflowUIPlugin: registered media routes (delete, view, list, tree, capabilities, version_info)")