from fastapi import FastAPI, Request
from scalar_fastapi import get_scalar_api_reference
import sentry_sdk
from sentry_sdk.integrations.loguru import LoguruIntegration, LoggingLevels
from sentry_sdk.integrations.fastapi import FastApiIntegration
from fastapi.middleware.cors import CORSMiddleware

from .utils.KVCache import MediaSourceKVCache
from .router import ffmpeg, cache, google, task, jm_router, mj_router, hl_router
from .models.settings.cluster import WorkerConfig

config = WorkerConfig()

web_app = FastAPI(title="Modal worker API",
                  version=config.api_version,
                  summary="Modal Worker的API, 包括缓存视频, 发起生产任务等",
                  servers=[
                      {
                          'url': f'https://bowongai-{config.modal_environment}--{config.modal_app_name}-fastapi-webapp.modal.run',
                          'description': '当前Modal API接口endpoint'
                      }
                  ])

web_app_tikhub = FastAPI(title="Modal Tikhub API",
                         version="1.0.0",
                         summary="Modal Tikhub API, 包含抖音Web API",
                         )

sentry_sdk.init(dsn="https://dab7b7ae652216282c89f029a76bb10a@sentry.bowongai.com/2",
                send_default_pii=True,
                traces_sample_rate=1.0,
                profiles_sample_rate=1.0,
                add_full_stack=True,
                environment=config.modal_environment,
                integrations=[
                    LoguruIntegration(level=LoggingLevels.INFO.value, event_level=LoggingLevels.ERROR.value),
                    FastApiIntegration()
                ]
                )
modal_kv_cache = MediaSourceKVCache(kv_name=config.modal_kv_name,
                                    environment=config.modal_environment, )

sentry_header_schema = {
    "x-trace-id": {
        "description": "Sentry Transaction ID",
        "schema": {
            "type": "string",
        }
    },
    "x-baggage": {
        "description": "Sentry Transaction baggage",
        "schema": {
            "type": "string",
        }
    }
}

ALIAS_MAP = {
    "/comfyui": "/comfyui/v2",
}

web_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@web_app.middleware("http")
@web_app.middleware("https")
async def alias_middleware(request: Request, call_next):
    if request.url.path in ALIAS_MAP.keys():
        request.scope["path"] = ALIAS_MAP[request.url.path]
    return await call_next(request)


@web_app.get("/scalar", include_in_schema=False)
async def scalar():
    return get_scalar_api_reference(openapi_url='/openapi.json', title="Modal worker web endpoint")


web_app.include_router(ffmpeg.router)
web_app.include_router(cache.router)

# web_app.include_router(comfyui.router)

web_app.include_router(google.router)
web_app.include_router(task.router)
web_app.include_router(jm_router.router)
web_app.include_router(mj_router.router)
web_app.include_router(hl_router.router)