import psutil
import time
import asyncio
from datetime import datetime
from collections import deque
from typing import Dict, Optional, List


from ..models.traffic_models import TrafficRecord
from src.backend.app.api.database import SessionLocal
from ..utils.logger import logger


class TrafficMonitor:
    def __init__(self, history_size: int = 300):
        self.history_size = history_size
        self.history = {
            "sent": deque(maxlen=history_size),
            "recv": deque(maxlen=history_size),
            "time": deque(maxlen=history_size),
            "interfaces": {}
        }
        self.running = False
        self.task = None
        self.update_interval = 1.0  # 秒

    @staticmethod
    def get_interfaces() -> List[str]:
        """获取所有网络接口名称"""
        return list(psutil.net_io_counters(pernic=True).keys())

    def start_monitoring(self):
        """启动流量监控"""
        if not self.running:
            self.running = True
            self.task = asyncio.create_task(self._monitor_loop())
            logger.info("流量监控已启动")

    async def stop_monitoring(self):
        """停止流量监控"""
        if self.running:
            self.running = False
            self.task.cancel()
            try:
                await self.task
            except asyncio.CancelledError:
                pass
            logger.info("流量监控已停止")

    async def _monitor_loop(self):
        """监控主循环"""
        last_stats = psutil.net_io_counters(pernic=True)
        last_time = time.time()

        while self.running:
            await asyncio.sleep(self.update_interval)

            current_time = time.time()
            current_stats = psutil.net_io_counters(pernic=True)
            elapsed = current_time - last_time

            for iface in current_stats:
                if iface not in self.history["interfaces"]:

                    self.history["interfaces"][iface] = {
                        "sent": deque(maxlen=self.history_size),
                        "recv": deque(maxlen=self.history_size)
                    }

                if iface in last_stats:
                    sent_rate = (current_stats[iface].bytes_sent - last_stats[iface].bytes_sent) / elapsed
                    recv_rate = (current_stats[iface].bytes_recv - last_stats[iface].bytes_recv) / elapsed


                    self.history["sent"].append(sent_rate)
                    self.history["recv"].append(recv_rate)
                    self.history["time"].append(datetime.now())


                    self.history["interfaces"][iface]["sent"].append(sent_rate)
                    self.history["interfaces"][iface]["recv"].append(recv_rate)


            self._save_to_db(current_stats)

            last_stats = current_stats
            last_time = current_time

    @staticmethod
    def _save_to_db(stats):
        """保存流量数据到数据库"""
        with SessionLocal() as session:
            for iface, counters in stats.items():
                record = TrafficRecord(
                    interface=iface,
                    bytes_sent=counters.bytes_sent,
                    bytes_recv=counters.bytes_recv,
                    packets_sent=counters.packets_sent,
                    packets_recv=counters.packets_recv,
                    timestamp=datetime.now()
                )
                session.add(record)
            session.commit()

    def get_current_traffic(self, interface: Optional[str] = None) -> Dict:
        """获取当前流量数据"""
        stats = psutil.net_io_counters(pernic=True)

        if interface:
            if interface in stats:
                return self._format_interface_stats(stats[interface])
            return {}

        return {iface: self._format_interface_stats(data) for iface, data in stats.items()}

    @staticmethod
    def _format_interface_stats(counters) -> Dict:
        """格式化接口统计数据"""
        return {
            "bytes_sent": counters.bytes_sent,
            "bytes_recv": counters.bytes_recv,
            "packets_sent": counters.packets_sent,
            "packets_recv": counters.packets_recv,
            "errin": counters.errin,
            "errout": counters.errout,
            "dropin": counters.dropin,
            "dropout": counters.dropout
        }

    def get_traffic_history(self, interface: Optional[str] = None) -> Dict:
        """获取流量历史数据"""
        if interface and interface in self.history["interfaces"]:
            return {
                "sent": list(self.history["interfaces"][interface]["sent"]),
                "recv": list(self.history["interfaces"][interface]["recv"]),
                "time": list(self.history["time"])
            }

        return {
            "sent": list(self.history["sent"]),
            "recv": list(self.history["recv"]),
            "time": list(self.history["time"])
        }

traffic_monitor = TrafficMonitor()
