from fastapi import Request
from sqlalchemy import Column, Integer, String
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# 创建数据库引擎
DATABASE_URL = "sqlite:///./api_count.db"
engine = create_engine(DATABASE_URL)

# 创建数据库会话
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建数据库模型基类
Base = declarative_base()


# 定义存储调用次数的数据库模型
class CallCount(Base):
    __tablename__ = "call_count"
    id = Column(Integer, primary_key=True, index=True)
    service_name = Column(String, nullable=False)
    call_count = Column(Integer, default=0)


# 创建数据库表（如果不存在）
Base.metadata.create_all(bind=engine)


# 定义中间件函数，用于记录接口调用次数

def count_api_calls(request: Request, call_next):
    db = SessionLocal()
    service_name = request.url.path
    call_record = db.query(CallCount).filter_by(service_name=service_name).first()
    if call_record is None:
        call_record = CallCount(service_name=service_name, call_count=1)
        db.add(call_record)
    else:
        call_record.call_count += 1
    db.commit()
    db.refresh(call_record)
    response = call_next(request)
    return response
