dashboard-nanobot/backend/api/image_router.py

from typing import Any, Dict, List

from fastapi import APIRouter, Depends, HTTPException
from sqlmodel import Session, select

from core.cache import cache
from core.database import get_session
from core.docker_instance import docker_manager
from models.bot import BotInstance, NanobotImage
from services.cache_service import _cache_key_images, _invalidate_images_cache

router = APIRouter()


def _serialize_image(row: NanobotImage) -> Dict[str, Any]:
    created_at = row.created_at.isoformat() + "Z" if row.created_at else None
    return {
        "tag": row.tag,
        "image_id": row.image_id,
        "version": row.version,
        "status": row.status,
        "source_dir": row.source_dir,
        "created_at": created_at,
    }


def _reconcile_registered_images(session: Session) -> None:
    rows = session.exec(select(NanobotImage)).all()
    dirty = False
    for row in rows:
        docker_exists = docker_manager.has_image(row.tag)
        next_status = "READY" if docker_exists else "ERROR"
        next_image_id = row.image_id
        if docker_exists and docker_manager.client:
            try:
                next_image_id = docker_manager.client.images.get(row.tag).id
            except Exception:
                next_image_id = row.image_id
        if row.status != next_status or row.image_id != next_image_id:
            row.status = next_status
            row.image_id = next_image_id
            session.add(row)
            dirty = True
    if dirty:
        session.commit()


@router.get("/api/images")
def list_images(session: Session = Depends(get_session)):
    cached = cache.get_json(_cache_key_images())
    if isinstance(cached, list) and all(isinstance(row, dict) for row in cached):
        return cached
    if isinstance(cached, list):
        _invalidate_images_cache()
    try:
        _reconcile_registered_images(session)
    except Exception as exc:
        # Docker status probing should not break image management in dev mode.
        print(f"[image_router] reconcile images skipped: {exc}")
    rows = session.exec(select(NanobotImage).order_by(NanobotImage.created_at.desc())).all()
    payload = [_serialize_image(row) for row in rows]
    cache.set_json(_cache_key_images(), payload, ttl=60)
    return payload

@router.delete("/api/images/{tag:path}")
def delete_image(tag: str, session: Session = Depends(get_session)):
    image = session.get(NanobotImage, tag)
    if not image:
        raise HTTPException(status_code=404, detail="Image not found")
    
    # 检查是否有机器人正在使用此镜像
    bots_using = session.exec(select(BotInstance).where(BotInstance.image_tag == tag)).all()
    if bots_using:
        raise HTTPException(status_code=400, detail=f"Cannot delete image: {len(bots_using)} bots are using it.")

    session.delete(image)
    session.commit()
    _invalidate_images_cache()
    return {"status": "deleted"}

@router.get("/api/docker-images")
def list_docker_images(repository: str = "nanobot-base"):
    rows = docker_manager.list_images_by_repo(repository)
    return rows

@router.post("/api/images/register")
def register_image(payload: dict, session: Session = Depends(get_session)):
    tag = (payload.get("tag") or "").strip()
    source_dir = (payload.get("source_dir") or "manual").strip() or "manual"
    if not tag:
        raise HTTPException(status_code=400, detail="tag is required")

    if not docker_manager.has_image(tag):
        raise HTTPException(status_code=404, detail=f"Docker image not found: {tag}")

    version = tag.split(":")[-1].removeprefix("v") if ":" in tag else tag
    try:
        docker_img = docker_manager.client.images.get(tag) if docker_manager.client else None
        image_id = docker_img.id if docker_img else None
    except Exception:
        image_id = None

    row = session.get(NanobotImage, tag)
    if not row:
        row = NanobotImage(
            tag=tag,
            version=version,
            status="READY",
            source_dir=source_dir,
            image_id=image_id,
        )
    else:
        row.version = version
        row.status = "READY"
        row.source_dir = source_dir
        row.image_id = image_id
    session.add(row)
    session.commit()
    session.refresh(row)
    _invalidate_images_cache()
    return _serialize_image(row)