from src.backend.app.api.database import Base  # 修复：导入 Base
from sqlalchemy import Column, Integer, String, DateTime, BigInteger, Float


class TrafficRecord(Base):
    """网络流量记录模型"""
    __tablename__ = "traffic_records"

    id = Column(Integer, primary_key=True, index=True)
    interface = Column(String(50), index=True)
    bytes_sent = Column(Integer)
    bytes_recv = Column(Integer)
    packets_sent = Column(Integer)
    packets_recv = Column(Integer)
    timestamp = Column(DateTime)

    def to_dict(self):
        return {
            "id": self.id,
            "interface": self.interface,
            "bytes_sent": self.bytes_sent,
            "bytes_recv": self.bytes_recv,
            "packets_sent": self.packets_sent,
            "packets_recv": self.packets_recv,
            "timestamp": self.timestamp.isoformat()
        }


class SwitchTrafficRecord(Base):
    __tablename__ = "switch_traffic_records"

    id = Column(Integer, primary_key=True, index=True)
    switch_ip = Column(String(50), index=True)
    interface = Column(String(50))
    bytes_in = Column(BigInteger)
    bytes_out = Column(BigInteger)
    rate_in = Column(Float)
    rate_out = Column(Float)
    timestamp = Column(DateTime)

    def to_dict(self):
        return {
            "id": self.id,
            "switch_ip": self.switch_ip,
            "interface": self.interface,
            "bytes_in": self.bytes_in,
            "bytes_out": self.bytes_out,
            "rate_in": self.rate_in,
            "rate_out": self.rate_out,
            "timestamp": self.timestamp.isoformat()
        }