Imago/app/routers/admin.py

from fastapi import APIRouter, Depends, HTTPException, Request
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from typing import Dict, Any, List
from pathlib import Path

from app.database import get_db
from app.models.image import Image
from app.models.client import APIClient
from app.dependencies.auth import require_scope

router = APIRouter(prefix="/admin/api", tags=["Admin"])

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
ALLOWED_DOCS = {
    "README.md": PROJECT_ROOT / "README.md",
    "USER_GUIDE.md": PROJECT_ROOT / "docs" / "USER_GUIDE.md",
    "ARCHITECTURE.md": PROJECT_ROOT / "docs" / "ARCHITECTURE.md",
    "API_GUIDE.md": PROJECT_ROOT / "docs" / "API_GUIDE.md",
    "SDK.md": PROJECT_ROOT / "docs" / "SDK.md",
    "SHAARLI-INTEGRATION.md": PROJECT_ROOT / "docs" / "SHAARLI-INTEGRATION.md",
    "WEBSOCKET.md": PROJECT_ROOT / "docs" / "WEBSOCKET.md",
}

@router.get("/stats")
async def get_stats(
    _: APIClient = Depends(require_scope("admin")),
    db: AsyncSession = Depends(get_db)
) -> Dict[str, Any]:
    """Get global system statistics."""
    # Count images
    total_images_query = await db.execute(select(func.count(Image.id)))
    total_images = total_images_query.scalar() or 0
    
    # Storage used (sum of file_size)
    storage_query = await db.execute(select(func.sum(Image.file_size)))
    storage_bytes = storage_query.scalar() or 0
    
    # Total AI tokens (sum of prompt + output tokens)
    tokens_query = await db.execute(
        select(func.sum(Image.ai_prompt_tokens) + func.sum(Image.ai_output_tokens))
        .where(
            Image.ai_prompt_tokens.isnot(None) | Image.ai_output_tokens.isnot(None)
        )
    )
    total_tokens = tokens_query.scalar() or 0

    # Total clients
    clients_query = await db.execute(select(func.count(APIClient.id)))
    total_clients = clients_query.scalar() or 0

    return {
        "total_images": total_images,
        "total_storage_bytes": storage_bytes,
        "total_ai_tokens": total_tokens,
        "total_clients": total_clients
    }

@router.get("/clients")
async def get_clients(
    _: APIClient = Depends(require_scope("admin")),
    db: AsyncSession = Depends(get_db)
) -> List[Dict[str, Any]]:
    """List all clients with metrics."""
    result = await db.execute(select(APIClient).order_by(APIClient.created_at.desc()))
    clients = result.scalars().all()
    
    return [
        {
            "id": c.id,
            "name": c.name,
            "is_active": c.is_active,
            "is_premium": c.plan == "premium",
            "storage_used": c.storage_used_bytes,
            "storage_quota": c.quota_storage_mb,
            "created_at": c.created_at
        } for c in clients
    ]

@router.get("/queue/status")
async def get_queue_status(
    request: Request,
    _: APIClient = Depends(require_scope("admin"))
) -> Dict[str, Any]:
    """Get ARQ queue status using the redis connection."""
    redis = request.app.state.redis
    
    # Note: A real implementation would parse the ARQ keys here. 
    # For now, we return basic statistics by probing redis directly with arq known queues.
    # Count pending jobs in arq:queue
    pending_count = await redis.llen("arq:queue") if hasattr(redis, "llen") else 0
    # In ARQ, active jobs are harder to count without querying the worker sets, 
    # so we return pending count.
    
    return {
        "pending_jobs": pending_count,
        "status": "active"
    }

@router.post("/clients/{client_id}/toggle")
async def toggle_client(
    client_id: str,
    _: APIClient = Depends(require_scope("admin")),
    db: AsyncSession = Depends(get_db)
) -> Dict[str, Any]:
    """Activate or deactivate a client."""
    result = await db.execute(select(APIClient).where(APIClient.id == client_id))
    client = result.scalar_one_or_none()
    if not client:
        raise HTTPException(status_code=404, detail="Client not found")
        
    client.is_active = not client.is_active
    await db.commit()
    
    return {"id": client.id, "is_active": client.is_active}

@router.post("/clients/{client_id}/reset-quota")
async def reset_client_quota(
    client_id: str,
    _: APIClient = Depends(require_scope("admin")),
    db: AsyncSession = Depends(get_db)
) -> Dict[str, Any]:
    """Reset the quota counters for a client."""
    result = await db.execute(select(APIClient).where(APIClient.id == client_id))
    client = result.scalar_one_or_none()
    if not client:
        raise HTTPException(status_code=404, detail="Client not found")
        
    client.storage_used_bytes = 0
    await db.commit()
    
    return {"id": client.id, "storage_used": 0}

@router.get("/docs")
async def list_docs(
    _: APIClient = Depends(require_scope("admin"))
) -> List[Dict[str, str]]:
    """List available documentation files."""
    docs = []
    for name, path in ALLOWED_DOCS.items():
        if path.exists():
            docs.append({"name": name, "title": name.replace(".md", "").replace("-", " ").title()})
        else:
            docs.append({"name": name, "title": name.replace(".md", "").replace("-", " ").title()}) # we might still want to list them, or omit
    return docs

@router.get("/docs/{filename}")
async def get_doc(
    filename: str,
    _: APIClient = Depends(require_scope("admin"))
) -> Dict[str, str]:
    """Get the content of a specific documentation file."""
    if filename not in ALLOWED_DOCS:
        raise HTTPException(status_code=404, detail="Document not found or access denied")
    
    path = ALLOWED_DOCS[filename]
    if not path.exists():
        raise HTTPException(status_code=404, detail=f"Document file {filename} not found on server")
        
    try:
        content = path.read_text(encoding="utf-8")
        return {"filename": filename, "content": content}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error reading document: {str(e)}")