import argparse
import gc
import os
import sys

import psutil
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from loguru import logger
from starlette.responses import JSONResponse

from routers import router

app = FastAPI(title="layoutlmv3_service")
app.state.request_counter = 0

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

app.include_router(router)

# ── 로거 설정 ─────────────────────────────────────
logger.remove()
logger.add(
    "logs/{time:YYYY-MM-DD}.log",
    level="DEBUG",
    rotation="00:00",
    retention="30 days",
    format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
)


# ── 메모리 유틸 ───────────────────────────────────
def log_memory_usage() -> float:
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024  # MB


# ── HTTP 미들웨어 ─────────────────────────────────
@app.middleware("http")
async def log_requests(request: Request, call_next):
    try:
        memory_before   = log_memory_usage()
        request_counter = getattr(app.state, "request_counter", 0) + 1
        app.state.request_counter = request_counter

        logger.info(f"request_counter: {request_counter}")
        logger.info(f"Received: {request.method} {request.url}")
        logger.info(f"Headers: {dict(request.headers)}")

        response = await call_next(request)
        logger.info(f"Response status: {response.status_code}")

        if request_counter % 100 == 0:
            logger.info("gc.collect 실행")
            gc.collect()

        memory_after = log_memory_usage()
        logger.info(
            f"Memory: Before={memory_before:.2f}MB, "
            f"After={memory_after:.2f}MB, "
            f"Diff={memory_after - memory_before:.2f}MB"
        )
        return response

    except Exception as e:
        logger.exception(f"요청 처리 중 예외 발생: {e}")
        return JSONResponse(
            status_code=500,
            content={"detail": "Internal Server Error", "error": str(e)},
        )


@app.get("/", tags=["기타"])
async def root():
    return {"message": "layoutlmv3 invoice classifier"}


@app.get("/health", tags=["기타"])
async def health():
    import config
    from utils.cache import load_cache
    return {
        "status":      "ok",
        "device":      str(config.DEVICE),
        "model_ready": (config.MODEL_SAVE_PATH / "config.json").exists(),
        "dataset_size": len(load_cache()),
    }


# ── 진입점 ────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-host",      default="0.0.0.0")
    parser.add_argument("-port",      default=10941)
    parser.add_argument("-pool_size", default="10")
    args = parser.parse_args()

    logger.info(f"서버 시작: {args.host}:{args.port}")
    uvicorn.run(
        app,
        host=args.host,
        port=int(args.port),
        limit_concurrency=1000,
        timeout_keep_alive=120,
        log_level="debug",
    )