import uvicorn
from fastapi import FastAPI
from starlette.staticfiles import StaticFiles
from Util.PostgreSQLUtil import init_postgres_pool, close_postgres_pool

from Routes.Ggb import router as ggb_router
from Routes.Knowledge import router as knowledge_router
from Routes.Oss import router as oss_router
from Routes.Rag import router as rag_router
from Routes.ZuoWen import router as zuowen_router
from Routes.Llm import router as llm_router
from Util.LightRagUtil import *
from contextlib import asynccontextmanager

# 控制日志输出
logger = logging.getLogger('lightrag')
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)


@asynccontextmanager
async def lifespan(_: FastAPI): 
    pool = await init_postgres_pool()
    app.state.pool = pool
    
    try:
        yield 
    finally:
        # 应用关闭时销毁连接池
        await close_postgres_pool(pool)
app = FastAPI(lifespan=lifespan)

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="Static"), name="static")

# 加载路由
app.include_router(zuowen_router)  # 作文批阅路由
app.include_router(ggb_router)  # Geogebra路由
app.include_router(rag_router)  # LightRAG路由
app.include_router(knowledge_router)  # 知识图谱路由
app.include_router(oss_router)  # 阿里云OSS路由

app.include_router(llm_router)  # 大模型路由

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8100)