from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from contextlib import asynccontextmanager

#from router import llamindex_router
from router import cnn_router

@asynccontextmanager
async def lifespan_manager(app: FastAPI):
    """
    서버 시작 시 모델을 로드하고 종료 시 정리합니다.
    """
    # 1. CNN 모델 로드 (cnn_router에 위임)
    await cnn_router.load_models()
    
    # 서버가 요청 처리를 시작하도록 제어권을 넘겨줍니다.
    yield 
    
    # 서버 종료 시 실행되는 정리(shutdown) 로직
    cnn_router.shutdown_models()

# FastAPI 애플리케이션 초기화
app = FastAPI(
    title="EfficientNetB0 이미지 분류 API",
    description="Fine-tuned EfficientNetB0 모델을 사용하여 이미지를 예측합니다.",
    lifespan=lifespan_manager 
)

# CORS 설정 (모든 도메인 허용)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

#app.include_router(llamindex_router.router, prefix="/llama_index")
app.include_router(cnn_router.router) # CNN 라우터 추가

# 헬스 체크용 기본 엔드포인트
@app.get("/", summary="API 헬스 체크")
def read_root():
    """API 서버가 정상적으로 작동하는지 확인합니다."""
    return {"message": "EfficientNetB0 Classification API is running successfully."}

if __name__ == "__main__":
    # --reload 옵션을 추가하여 코드가 변경될 때마다 자동 재시작되게 설정합니다.
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)