import uvicorn
from fastapi import FastAPI
from contextlib import asynccontextmanager
import argparse

from pt_gen.core import config
from pt_gen.api import endpoints
from pt_gen.api.endpoints import get_orchestrator


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 应用启动时
    print("应用启动...")
    # 预热依赖：第一次调用 get_orchestrator 来创建实例并建立 redis 连接
    orchestrator = get_orchestrator()
    print("Orchestrator 实例已创建。")
    
    yield
    # 应用关闭时
    print("应用关闭，关闭 Redis 连接...")
    # 获取实例并关闭连接
    orchestrator = get_orchestrator()
    await orchestrator.cache.close()
    print("Redis 连接已关闭。")

app = FastAPI(title="PT-Gen API", lifespan=lifespan)
app.include_router(endpoints.router, prefix="/api", tags=["Generator"])

@app.get("/", tags=["Root"])
def read_root():
    return {"message": "Welcome to PT-Gen API"}


if __name__ == "__main__":
    # 解析命令行参数来获取 host 和 port
    parser = argparse.ArgumentParser(description="Run PT-Gen API server.")
    parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to bind")
    parser.add_argument("--port", type=int, default=8000, help="Port to listen on")
    parser.add_argument("--config", type=str, default="configs/config.yaml", help="Path to the config file")
    
    cli_args = parser.parse_args()
    
    uvicorn.run(
        "main:app",
        host=cli_args.host,
        port=cli_args.port,
        reload=True # 开发时开启，生产环境应为 False
    )