from datetime import datetime, timedelta
from fastapi import (APIRouter, HTTPException, Response, WebSocket, WebSocketDisconnect)
from typing import List
from pydantic import BaseModel
import asyncio
from fastapi.responses import HTMLResponse
import matplotlib.pyplot as plt
import io
import base64

from ..services.switch_traffic_monitor import get_switch_monitor
from ..utils import logger
from ...app.services.ai_services import AIService
from ...app.api.network_config import SwitchConfigurator
from ...config import settings
from ..services.network_scanner import NetworkScanner
from ...app.services.traffic_monitor import traffic_monitor
from ...app.models.traffic_models import TrafficRecord, SwitchTrafficRecord
from src.backend.app.api.database import SessionLocal




router = APIRouter(prefix="", tags=["API"])
scanner = NetworkScanner()

# 添加根路径路由
@router.get("/", include_in_schema=False)
async def root():
    return {
        "message": "欢迎使用AI交换机配置系统",
        "docs": f"{settings.API_PREFIX}/docs",
        "redoc": f"{settings.API_PREFIX}/redoc",
        "endpoints": [
            "/parse_command",
            "/apply_config",
            "/scan_network",
            "/list_devices",
            "/batch_apply_config"
            "/traffic/switch/current",  # 添加交换机流量端点
            "/traffic/switch/history"  # 添加交换机历史流量端点
        ]
    }

# 添加favicon处理
@router.get("/favicon.ico", include_in_schema=False)
async def favicon():
    return Response(status_code=204)

class BatchConfigRequest(BaseModel):
    config: dict
    switch_ips: List[str]  # 支持多个IP

@router.post("/batch_apply_config")
async def batch_apply_config(request: BatchConfigRequest):
    results = {}
    for ip in request.switch_ips:
        try:
            configurator = SwitchConfigurator()
            results[ip] = await configurator.apply_config(ip, request.config)
        except Exception as e:
            results[ip] = str(e)
    return {"results": results}



@router.get("/test")
async def test_endpoint():
    return {"message": "Hello World"}

@router.get("/scan_network", summary="扫描网络中的交换机")
async def scan_network(subnet: str = "192.168.1.0/24"):
    try:
        devices = scanner.scan_subnet(subnet)
        return {
            "success": True,
            "devices": devices,
            "count": len(devices)
        }
    except Exception as e:
        raise HTTPException(500, f"扫描失败: {str(e)}")

@router.get("/list_devices", summary="列出已发现的交换机")
async def list_devices():
    return {
        "devices": scanner.load_cached_devices()
    }

class CommandRequest(BaseModel):
    command: str

class ConfigRequest(BaseModel):
    config: dict
    switch_ip: str

@router.post("/parse_command", response_model=dict)
async def parse_command(request: CommandRequest):
    """
    解析中文命令并返回JSON配置
    """
    try:
        ai_service = AIService(settings.SILICONFLOW_API_KEY, settings.SILICONFLOW_API_URL)
        config = await ai_service.parse_command(request.command)
        return {"success": True, "config": config}
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=f"Failed to parse command: {str(e)}"
        )

@router.post("/apply_config", response_model=dict)
async def apply_config(request: ConfigRequest):
    """
    应用配置到交换机
    """
    try:
        configurator = SwitchConfigurator(
            username=settings.SWITCH_USERNAME,
            password=settings.SWITCH_PASSWORD,
            timeout=settings.SWITCH_TIMEOUT
        )
        result = await configurator.safe_apply(request.switch_ip, request.config)
        return {"success": True, "result": result}
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to apply config: {str(e)}"
        )
@router.get("/traffic/interfaces", summary="获取所有网络接口")
async def get_network_interfaces():
    return {
        "interfaces": traffic_monitor.get_interfaces()
    }

@router.get("/traffic/current", summary="获取当前流量数据")
async def get_current_traffic(interface: str = None):
    return traffic_monitor.get_current_traffic(interface)

@router.get("/traffic/history", summary="获取流量历史数据")
async def get_traffic_history(interface: str = None, limit: int = 100):
    history = traffic_monitor.get_traffic_history(interface)
    return {
        "sent": history["sent"][-limit:],
        "recv": history["recv"][-limit:],
        "time": [t.isoformat() for t in history["time"]][-limit:]
    }

@router.get("/traffic/records", summary="获取流量记录")
async def get_traffic_records(interface: str = None, limit: int = 100):
    with SessionLocal() as session:
        query = session.query(TrafficRecord)
        if interface:
            query = query.filter(TrafficRecord.interface == interface)
        records = query.order_by(TrafficRecord.timestamp.desc()).limit(limit).all()
        return [record.to_dict() for record in records]

@router.websocket("/ws/traffic")
async def websocket_traffic(websocket: WebSocket):
    """实时流量WebSocket"""
    await websocket.accept()
    try:
        while True:
            # 获取所有接口的当前流量
            traffic_data = traffic_monitor.get_current_traffic()
            await websocket.send_json(traffic_data)
            await asyncio.sleep(1)  # 每秒更新一次
    except WebSocketDisconnect:
        print("客户端断开连接")


@router.get("/", include_in_schema=False)
async def root():
    return {
        "message": "欢迎使用AI交换机配置系统",
        "docs": f"{settings.API_PREFIX}/docs",
        "redoc": f"{settings.API_PREFIX}/redoc",
        "endpoints": [
            "/parse_command",
            "/apply_config",
            "/scan_network",
            "/list_devices",
            "/batch_apply_config",
            "/traffic/switch/current",  # 添加交换机流量端点
            "/traffic/switch/history"  # 添加交换机历史流量端点
        ]
    }


# ... 其他路由保持不变 ...

@router.get("/traffic/switch/interfaces", summary="获取交换机的网络接口")
async def get_switch_interfaces(switch_ip: str):
    """获取指定交换机的所有接口"""
    try:
        monitor = get_switch_monitor(switch_ip)
        interfaces = list(monitor.interface_oids.keys())
        return {
            "switch_ip": switch_ip,
            "interfaces": interfaces
        }
    except Exception as e:
        logger.error(f"获取交换机接口失败: {str(e)}")
        raise HTTPException(500, f"获取接口失败: {str(e)}")


@router.get("/traffic/switch/current", summary="获取交换机的当前流量数据")
async def get_switch_current_traffic(switch_ip: str, interface: str = None):
    """获取交换机的当前流量数据"""
    try:
        monitor = get_switch_monitor(switch_ip)

        # 如果没有指定接口，获取所有接口的流量
        if not interface:
            traffic_data = {}
            for iface in monitor.interface_oids:
                traffic_data[iface] = await get_interface_current_traffic(switch_ip, iface)
            return {
                "switch_ip": switch_ip,
                "traffic": traffic_data
            }

        # 获取指定接口的流量
        return await get_interface_current_traffic(switch_ip, interface)
    except Exception as e:
        logger.error(f"获取交换机流量失败: {str(e)}")
        raise HTTPException(500, f"获取流量失败: {str(e)}")


async def get_interface_current_traffic(switch_ip: str, interface: str) -> dict:
    """获取指定交换机接口的当前流量数据"""
    try:
        with SessionLocal() as session:
            # 获取最新记录
            record = session.query(SwitchTrafficRecord).filter(
                SwitchTrafficRecord.switch_ip == switch_ip,
                SwitchTrafficRecord.interface == interface
            ).order_by(SwitchTrafficRecord.timestamp.desc()).first()

            if not record:
                return {
                    "switch_ip": switch_ip,
                    "interface": interface,
                    "rate_in": 0.0,
                    "rate_out": 0.0,
                    "bytes_in": 0,
                    "bytes_out": 0
                }

            return {
                "switch_ip": switch_ip,
                "interface": interface,
                "rate_in": record.rate_in,
                "rate_out": record.rate_out,
                "bytes_in": record.bytes_in,
                "bytes_out": record.bytes_out
            }
    except Exception as e:
        logger.error(f"获取接口流量失败: {str(e)}")
        raise HTTPException(500, f"获取接口流量失败: {str(e)}")


@router.get("/traffic/switch/history", summary="获取交换机的流量历史数据")
async def get_switch_traffic_history(switch_ip: str, interface: str = None, minutes: int = 10):
    """获取交换机的流量历史数据"""
    try:
        monitor = get_switch_monitor(switch_ip)

        # 如果没有指定接口，返回所有接口的历史数据
        if not interface:
            return {
                "switch_ip": switch_ip,
                "history": monitor.get_traffic_history()
            }

        # 获取指定接口的历史数据
        with SessionLocal() as session:
            # 计算时间范围
            time_threshold = datetime.now() - timedelta(minutes=minutes)

            # 查询历史记录
            records = session.query(SwitchTrafficRecord).filter(
                SwitchTrafficRecord.switch_ip == switch_ip,
                SwitchTrafficRecord.interface == interface,
                SwitchTrafficRecord.timestamp >= time_threshold
            ).order_by(SwitchTrafficRecord.timestamp.asc()).all()

            # 提取数据
            history_data = {
                "in": [record.rate_in for record in records],
                "out": [record.rate_out for record in records],
                "time": [record.timestamp.isoformat() for record in records]
            }

            return {
                "switch_ip": switch_ip,
                "interface": interface,
                "history": history_data
            }
    except Exception as e:
        logger.error(f"获取历史流量失败: {str(e)}")
        raise HTTPException(500, f"获取历史流量失败: {str(e)}")


@router.websocket("/ws/traffic/switch")
async def websocket_switch_traffic(websocket: WebSocket, switch_ip: str, interface: str = None):
    """交换机实时流量WebSocket"""
    await websocket.accept()
    try:
        monitor = get_switch_monitor(switch_ip)

        while True:
            # 获取流量数据
            if interface:
                # 单个接口
                traffic_data = await get_interface_current_traffic(switch_ip, interface)
                await websocket.send_json(traffic_data)
            else:
                # 所有接口
                traffic_data = {}
                for iface in monitor.interface_oids:
                    traffic_data[iface] = await get_interface_current_traffic(switch_ip, iface)

                await websocket.send_json({
                    "switch_ip": switch_ip,
                    "traffic": traffic_data
                })

            # 每秒更新一次
            await asyncio.sleep(1)
    except WebSocketDisconnect:
        logger.info(f"客户端断开交换机流量连接: {switch_ip}")
    except Exception as e:
        logger.error(f"交换机流量WebSocket错误: {str(e)}")
        await websocket.close(code=1011, reason=str(e))


# 交换机流量可视化
@router.get("/traffic/switch/plot", response_class=HTMLResponse, summary="交换机流量可视化")
async def plot_switch_traffic(switch_ip: str, interface: str, minutes: int = 10):
    """生成交换机流量图表"""
    try:
        # 获取历史数据
        history = await get_switch_traffic_history(switch_ip, interface, minutes)
        history_data = history["history"]

        # 提取数据点
        time_points = [datetime.fromisoformat(t) for t in history_data["time"]]
        in_rates = history_data["in"]
        out_rates = history_data["out"]

        # 创建图表
        plt.figure(figsize=(12, 6))
        plt.plot(time_points, in_rates, label="流入流量 (B/s)")
        plt.plot(time_points, out_rates, label="流出流量 (B/s)")
        plt.title(f"交换机 {switch_ip} 接口 {interface} 流量监控 - 最近 {minutes} 分钟")
        plt.xlabel("时间")
        plt.ylabel("流量 (字节/秒)")
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        plt.tight_layout()

        # 转换为HTML图像
        buf = io.BytesIO()
        plt.savefig(buf, format="png")
        buf.seek(0)
        image_base64 = base64.b64encode(buf.read()).decode("utf-8")
        plt.close()

        return f"""
        <html>
            <head>
                <title>交换机流量监控</title>
                <style>
                    body {{ font-family: Arial, sans-serif; margin: 20px; }}
                    .container {{ max-width: 900px; margin: 0 auto; }}
                </style>
            </head>
            <body>
                <div class="container">
                    <h1>交换机 {switch_ip} 接口 {interface} 流量监控</h1>
                    <img src="data:image/png;base64,{image_base64}" alt="流量图表">
                    <p>更新时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
                </div>
            </body>
        </html>
        """
    except Exception as e:
        logger.error(f"生成流量图表失败: {str(e)}")
        return HTMLResponse(content=f"<h1>错误</h1><p>{str(e)}</p>", status_code=500)