test/moxie/api/admin.py

"""
Hidden Admin UI Routes
Configuration, Document Upload, and ComfyUI Workflow Management
"""
import json
import os
from pathlib import Path
from typing import Optional
from fastapi import APIRouter, Request, UploadFile, File, Form, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from loguru import logger

from config import settings, get_data_dir, get_workflows_dir, save_config_to_db, load_config_from_db


router = APIRouter()

# Templates
templates = Jinja2Templates(directory=Path(__file__).parent.parent / "admin" / "templates")


# ============================================================================
# Config Models
# ============================================================================

class EndpointConfig(BaseModel):
    """API endpoint configuration."""
    gemini_api_key: Optional[str] = None
    gemini_model: str = "gemini-1.5-flash"
    openrouter_api_key: Optional[str] = None
    openrouter_model: str = "meta-llama/llama-3-8b-instruct:free"
    comfyui_host: str = "http://127.0.0.1:8188"
    ollama_host: str = "http://127.0.0.1:11434"
    ollama_model: str = "qwen2.5:2b"
    embedding_model: str = "qwen3-embedding:4b"


# ============================================================================
# Admin UI Routes
# ============================================================================

@router.get("/", response_class=HTMLResponse)
async def admin_dashboard(request: Request):
    """Admin dashboard homepage."""
    config = load_config_from_db()
    return templates.TemplateResponse(
        "dashboard.html",
        {
            "request": request,
            "config": config,
            "settings": settings
        }
    )


@router.get("/endpoints", response_class=HTMLResponse)
async def endpoints_page(request: Request):
    """API endpoint configuration page."""
    config = load_config_from_db()
    return templates.TemplateResponse(
        "endpoints.html",
        {
            "request": request,
            "config": config,
            "settings": settings
        }
    )


@router.post("/endpoints")
async def save_endpoints(config: EndpointConfig):
    """Save endpoint configuration to database."""
    config_dict = config.model_dump(exclude_none=True)
    for key, value in config_dict.items():
        save_config_to_db(key, value)
    
    logger.info("Endpoint configuration saved")
    return {"status": "success", "message": "Configuration saved"}


@router.get("/documents", response_class=HTMLResponse)
async def documents_page(request: Request):
    """Document management page."""
    rag_store = request.app.state.rag_store
    
    documents = rag_store.list_documents()
    return templates.TemplateResponse(
        "documents.html",
        {
            "request": request,
            "documents": documents,
            "settings": settings
        }
    )


@router.post("/documents/upload")
async def upload_document(
    request: Request,
    file: UploadFile = File(...),
    chunk_size: int = Form(default=500),
    overlap: int = Form(default=50)
):
    """Upload and index a document."""
    rag_store = request.app.state.rag_store
    
    # Read file content
    content = await file.read()
    
    # Process based on file type
    filename = file.filename or "unknown"
    file_ext = Path(filename).suffix.lower()
    
    try:
        doc_id = await rag_store.add_document(
            filename=filename,
            content=content,
            file_type=file_ext,
            chunk_size=chunk_size,
            overlap=overlap
        )
        
        logger.info(f"Document uploaded: {filename} (ID: {doc_id})")
        return {"status": "success", "document_id": doc_id, "filename": filename}
    
    except Exception as e:
        logger.error(f"Failed to upload document: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.delete("/documents/{doc_id}")
async def delete_document(doc_id: str, request: Request):
    """Delete a document from the store."""
    rag_store = request.app.state.rag_store
    
    try:
        rag_store.delete_document(doc_id)
        logger.info(f"Document deleted: {doc_id}")
        return {"status": "success"}
    except Exception as e:
        logger.error(f"Failed to delete document: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/comfyui", response_class=HTMLResponse)
async def comfyui_page(request: Request):
    """ComfyUI workflow management page."""
    config = load_config_from_db()
    workflows_dir = get_workflows_dir()
    
    workflows = {
        "image": None,
        "video": None,
        "audio": None
    }
    
    for workflow_type in workflows.keys():
        workflow_path = workflows_dir / f"{workflow_type}.json"
        if workflow_path.exists():
            with open(workflow_path, "r") as f:
                workflows[workflow_type] = json.load(f)
    
    return templates.TemplateResponse(
        "comfyui.html",
        {
            "request": request,
            "config": config,
            "workflows": workflows,
            "workflows_dir": str(workflows_dir),
            "settings": settings
        }
    )


@router.post("/comfyui/upload")
async def upload_comfyui_workflow(
    workflow_type: str = Form(...),
    file: UploadFile = File(...)
):
    """Upload a ComfyUI workflow JSON file."""
    if workflow_type not in ["image", "video", "audio"]:
        raise HTTPException(status_code=400, detail="Invalid workflow type")
    
    workflows_dir = get_workflows_dir()
    workflow_path = workflows_dir / f"{workflow_type}.json"
    
    try:
        content = await file.read()
        # Validate JSON
        workflow_data = json.loads(content)
        
        with open(workflow_path, "wb") as f:
            f.write(content)
        
        logger.info(f"ComfyUI workflow uploaded: {workflow_type}")
        return {"status": "success", "workflow_type": workflow_type}
    
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON file")
    except Exception as e:
        logger.error(f"Failed to upload workflow: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/comfyui/{workflow_type}")
async def get_comfyui_workflow(workflow_type: str):
    """Get a ComfyUI workflow JSON."""
    if workflow_type not in ["image", "video", "audio"]:
        raise HTTPException(status_code=400, detail="Invalid workflow type")
    
    workflows_dir = get_workflows_dir()
    workflow_path = workflows_dir / f"{workflow_type}.json"
    
    if not workflow_path.exists():
        raise HTTPException(status_code=404, detail="Workflow not found")
    
    with open(workflow_path, "r") as f:
        return json.load(f)


@router.delete("/comfyui/{workflow_type}")
async def delete_comfyui_workflow(workflow_type: str):
    """Delete a ComfyUI workflow."""
    if workflow_type not in ["image", "video", "audio"]:
        raise HTTPException(status_code=400, detail="Invalid workflow type")
    
    workflows_dir = get_workflows_dir()
    workflow_path = workflows_dir / f"{workflow_type}.json"
    
    if workflow_path.exists():
        workflow_path.unlink()
        logger.info(f"ComfyUI workflow deleted: {workflow_type}")
    
    return {"status": "success"}


@router.get("/status")
async def get_status(request: Request):
    """Get system status."""
    rag_store = request.app.state.rag_store
    config = load_config_from_db()
    
    # Check Ollama connectivity
    ollama_status = "unknown"
    try:
        import httpx
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"{config.get('ollama_host', settings.ollama_host)}/api/tags", timeout=5.0)
            ollama_status = "connected" if resp.status_code == 200 else "error"
    except Exception:
        ollama_status = "disconnected"
    
    # Check ComfyUI connectivity
    comfyui_status = "unknown"
    try:
        import httpx
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"{config.get('comfyui_host', settings.comfyui_host)}/system_stats", timeout=5.0)
            comfyui_status = "connected" if resp.status_code == 200 else "error"
    except Exception:
        comfyui_status = "disconnected"
    
    return {
        "ollama": ollama_status,
        "comfyui": comfyui_status,
        "documents_count": rag_store.get_document_count(),
        "chunks_count": rag_store.get_chunk_count(),
    }