from fastapi import APIRouter, Depends, HTTPException, status from sqlalchemy.orm import Session from typing import List from .. import models, schemas, auth from ..database import get_db from ..db_manager import DbManager from ..scheduler import sync_scheduler_config from fastapi.responses import FileResponse from fastapi import UploadFile, File router = APIRouter( prefix="/admin/db", tags=["Admin Database"] ) @router.get("/backups", response_model=List[schemas.BackupInfo]) def get_backups( db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """List available database backups.""" return DbManager.get_backup_list() @router.get("/stats", response_model=schemas.DatabaseStats) def get_db_stats( db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Get database backup storage statistics.""" return DbManager.get_stats() @router.post("/backup", response_model=schemas.BackupInfo) def trigger_manual_backup( db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Trigger a manual database backup.""" filename = DbManager.create_backup(db, label="manual", user_id=current_admin.sub) # Re-fetch the newly created file info backups = DbManager.get_backup_list() for b in backups: if b.filename == filename: return b raise HTTPException(status_code=500, detail="Backup created but info not found") @router.post("/restore") def restore_database( payload: dict, db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Restore database from a specific file. DANGEROUS.""" filename = payload.get("filename") confirm = payload.get("confirm", False) if not filename: raise HTTPException(status_code=400, detail="Filename required") if not confirm: raise HTTPException(status_code=400, detail="Confirmation required") try: success = DbManager.restore_backup(filename, db, user_id=current_admin.sub) return {"status": "success", "message": f"Database restored from {filename}"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @router.get("/settings", response_model=schemas.DbSettingsUpdate) def get_db_settings( db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Get database retention and scheduling settings.""" # Ensure default settings exist retention = db.query(models.SystemSetting).filter(models.SystemSetting.key == "backup_retention_count").first() hour = db.query(models.SystemSetting).filter(models.SystemSetting.key == "backup_schedule_hour").first() freq = db.query(models.SystemSetting).filter(models.SystemSetting.key == "backup_schedule_freq_days").first() return { "retention_count": int(retention.value) if retention else 10, "schedule_hour": int(hour.value) if hour else 3, "schedule_freq_days": int(freq.value) if freq else 1 } @router.patch("/settings", response_model=schemas.DbSettingsUpdate) def update_db_settings( settings: schemas.DbSettingsUpdate, db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Update database settings and re-trigger scheduler sync.""" pairs = { "backup_retention_count": str(settings.retention_count), "backup_schedule_hour": str(settings.schedule_hour), "backup_schedule_freq_days": str(settings.schedule_freq_days) } for key, val in pairs.items(): existing = db.query(models.SystemSetting).filter(models.SystemSetting.key == key).first() if existing: existing.value = val else: db.add(models.SystemSetting(key=key, value=val)) db.commit() # Re-trigger scheduler sync sync_scheduler_config() return settings @router.get("/export") def export_database( current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Download the current database file.""" try: path = DbManager.export_db() return FileResponse( path, media_type="application/x-sqlite3", filename="inventory_export.db" ) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @router.post("/import") async def import_database( file: UploadFile = File(...), db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Upload and replace the current database. DANGEROUS.""" contents = await file.read() try: DbManager.import_db(contents, db, user_id=current_admin.sub) return {"status": "success", "message": "Database successfully imported and replaced."} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) import os from ..database import BASE_DIR PROJECT_ROOT = os.path.dirname(BASE_DIR) PROMPT_FILE_PATH = os.path.join(PROJECT_ROOT, "config", "ai_prompt.md") @router.get("/settings/prompt") def get_ai_prompt( db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Get the current AI extraction prompt (prioritizes config file).""" # 1. Try file first if os.path.exists(PROMPT_FILE_PATH): try: with open(PROMPT_FILE_PATH, 'r', encoding='utf-8') as f: return {"value": f.read().strip(), "source": "file"} except Exception as e: print(f"Error reading prompt file: {e}") # 2. Fallback to DB setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first() if not setting: return {"value": "", "source": "none"} return {"value": setting.value, "source": "database"} @router.post("/settings/prompt") def update_ai_prompt( payload: dict, db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Update the AI extraction prompt (writes to both file and DB).""" value = payload.get("value") if value is None: raise HTTPException(status_code=400, detail="Value required") # 1. Update File (Primary) try: os.makedirs(os.path.dirname(PROMPT_FILE_PATH), exist_ok=True) with open(PROMPT_FILE_PATH, 'w', encoding='utf-8') as f: f.write(value) except Exception as e: print(f"Failed to write prompt file: {e}") # We continue to DB update anyway # 2. Update DB (Backup/Sync) existing = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_extraction_prompt").first() if existing: existing.value = value else: db.add(models.SystemSetting(key="ai_extraction_prompt", value=value)) db.commit() return {"status": "success", "file_updated": os.path.exists(PROMPT_FILE_PATH)} from ..config_manager import ConfigManager @router.get("/settings/ai") def get_ai_config( db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Check AI provider status and active provider.""" gemini_key = os.environ.get("GEMINI_API_KEY") claude_key = os.environ.get("CLAUDE_API_KEY") # Get active provider from DB provider_setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_provider").first() active_provider = provider_setting.value if provider_setting else "gemini" return { "active_provider": active_provider, "providers": [ { "id": "gemini", "name": "Google Gemini 2.0", "configured": bool(gemini_key), "active": active_provider == "gemini", "masked_key": ConfigManager.get_masked_key("GEMINI_API_KEY") }, { "id": "claude", "name": "Anthropic Claude 3.5", "configured": bool(claude_key), "active": active_provider == "claude", "masked_key": ConfigManager.get_masked_key("CLAUDE_API_KEY") } ] } @router.post("/settings/ai-keys") def update_ai_keys( payload: dict, current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Update AI API keys in inventory.env.""" gemini_key = payload.get("gemini_api_key") claude_key = payload.get("claude_api_key") updates = {} if gemini_key: updates["GEMINI_API_KEY"] = gemini_key if claude_key: updates["CLAUDE_API_KEY"] = claude_key if not updates: return {"status": "no_change"} ConfigManager.update_keys(updates) return { "status": "success", "gemini_configured": bool(os.environ.get("GEMINI_API_KEY")), "claude_configured": bool(os.environ.get("CLAUDE_API_KEY")) } @router.post("/settings/test-ai-key") def test_ai_key( payload: dict, current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Test AI API key connectivity.""" provider = payload.get("provider") key = payload.get("key") if not provider or provider not in ["gemini", "claude"]: raise HTTPException(status_code=400, detail="Invalid provider") # If key is masked or empty, use the one from environment if not key or "****" in key: key = os.environ.get("GEMINI_API_KEY" if provider == "gemini" else "CLAUDE_API_KEY") if not key: raise HTTPException(status_code=400, detail="No API key provided or configured") try: if provider == "gemini": from google import genai client = genai.Client(api_key=key, http_options={'api_version': 'v1beta'}) # Listing models is the standard way to check key validity client.models.list() return {"status": "success", "message": "Google Gemini API connection verified!"} elif provider == "claude": import anthropic client = anthropic.Anthropic(api_key=key) # Claude also supports model listing for key verification client.models.list(limit=1) return {"status": "success", "message": "Anthropic Claude API connection verified!"} except Exception as e: error_msg = str(e) if "API_KEY_INVALID" in error_msg or "invalid_api_key" in error_msg: error_msg = "The API key provided is invalid." elif "PERMISSION_DENIED" in error_msg: error_msg = "Permission denied. Check if the key has correct permissions." raise HTTPException(status_code=400, detail=f"{provider.capitalize()} Test Failed: {error_msg}") @router.post("/settings/ai") def update_ai_provider( payload: dict, db: Session = Depends(get_db), current_admin: auth.TokenData = Depends(auth.get_current_admin) ): """Update the active AI provider.""" provider = payload.get("provider") if provider not in ["gemini", "claude"]: raise HTTPException(status_code=400, detail="Invalid provider. Must be 'gemini' or 'claude'.") existing = db.query(models.SystemSetting).filter(models.SystemSetting.key == "ai_provider").first() if existing: existing.value = provider else: db.add(models.SystemSetting(key="ai_provider", value=provider)) db.commit() return {"status": "success", "active_provider": provider}