# Viewer "Extend" support: from a finished VIDEO, pull its LAST frame (staged in
# the input dir for a LoadImage to consume) plus its dimensions/fps, and stage a
# content-addressed copy of the source video so the re-render can be recorded as
# a lineage CHILD (parent_filename). The frontend then rebuilds an I2V
# continuation graph (last frame -> next clip) the same way Re-pose rebuilds from
# a recipe template — see js/lib/extend-from-image.js.

import hashlib
import io as _io
import os

import folder_paths
import server
from aiohttp import web

from .api_utils import error_response
from .history_db import resolve_image_path

routes = server.PromptServer.instance.routes

# Same prefix Upscale uses: LoadImage's combo only lists top-level input files,
# so the namespace lives in the filename and the copy sits in the input ROOT.
_SRC_PREFIX = "promptchain_source_"
_VIDEO_EXTS = {".mp4", ".webm", ".mkv", ".mov", ".m4v", ".avi"}


def _staged_copy(raw: bytes, ext: str) -> str:
    """Content-addressed copy in the input root; same bytes -> same name, written
    once. Returns the bare filename for LoadImage / parent_filename."""
    digest = hashlib.sha256(raw).hexdigest()[:12]
    name = f"{_SRC_PREFIX}{digest}{ext}"
    dest = os.path.join(folder_paths.get_input_directory(), name)
    if not os.path.exists(dest):
        with open(dest, "wb") as f:
            f.write(raw)
    return name


def _last_frame_and_meta(video_bytes: bytes):
    """(PIL.Image last_frame, width, height, fps, frame_count). Decodes the whole
    stream — Wan clips are short (~81-121 frames) so a full pass is cheap and far
    more reliable than seeking to EOF."""
    import av  # PyAV (av) ships in the ComfyUI venv; powers the poster-frame path too

    last = None
    fps = 16.0
    count = 0
    container = av.open(_io.BytesIO(video_bytes))
    try:
        stream = container.streams.video[0]
        try:
            if stream.average_rate:
                fps = float(stream.average_rate)
        except Exception:
            fps = 16.0
        for frame in container.decode(stream):
            last = frame
            count += 1
    finally:
        container.close()
    if last is None:
        return None, 0, 0, fps, 0
    img = last.to_image().convert("RGB")
    w, h = img.size
    return img, w, h, fps, count


@routes.get("/promptchain/extend/prepare")
async def _api_extend_prepare(request):
    """Stage the source video's last frame + a lineage copy. ?hash=<image hash>."""
    image_hash = request.query.get("hash", "").strip()
    if not image_hash:
        return error_response("need hash")
    path = resolve_image_path(image_hash)
    if not path or not os.path.isfile(str(path)):
        return error_response("video not found for hash", 404)
    path = str(path)
    if os.path.splitext(path)[1].lower() not in _VIDEO_EXTS:
        return error_response("not a video file")

    try:
        with open(path, "rb") as f:
            video_bytes = f.read()
    except OSError as e:
        return error_response(f"could not read video: {e}", 500)

    try:
        import av  # noqa: F401
    except Exception:
        return error_response("PyAV (av) not available — cannot read video frames", 500)

    try:
        source_ref = _staged_copy(video_bytes, os.path.splitext(path)[1].lower())
        img, width, height, fps, count = _last_frame_and_meta(video_bytes)
        if img is None:
            return error_response("could not decode any video frame", 500)
        buf = _io.BytesIO()
        img.save(buf, format="PNG")
        last_frame = _staged_copy(buf.getvalue(), ".png")
    except Exception as e:
        return error_response(f"extend prepare failed: {e}", 500)

    return web.json_response({
        "last_frame": last_frame,
        "source_ref": source_ref,
        "width": width,
        "height": height,
        "fps": round(fps, 3),
        "frame_count": count,
    })