from fastapi import APIRouter, HTTPException, Depends
import pymysql
from utils.database import get_db_connection, DatabaseConfig
from pydantic import BaseModel

router = APIRouter(
    prefix="/traffic-data",
    tags=["traffic-data"]
)

class TrafficDataRequest(BaseModel):
    speed: float
    var_speed: float
    waiting_time: float
    var_waiting: float
    avg_throughput: float
    avg_queue_length: float
    avg_occupancy: float
    reward: float
    last_action: int
    action: int

@router.post("/add")
def add_traffic_data(data: TrafficDataRequest):
    connection = get_db_connection(DatabaseConfig.TRAFFIC_MANAGER)
    try:
        with connection.cursor() as cursor:
            sql = """
                INSERT INTO traffic_data (
                    speed, var_speed, waiting_time, var_waiting, avg_throughput, avg_queue_length, avg_occupancy, reward, last_action, action
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            cursor.execute(sql, (
                data.speed,
                data.var_speed,
                data.waiting_time,
                data.var_waiting,
                data.avg_throughput,
                data.avg_queue_length,
                data.avg_occupancy,
                data.reward,
                data.last_action,
                data.action
            ))
            connection.commit()
            return {"message": "Traffic data added successfully"}
    except pymysql.Error as e:
        raise HTTPException(status_code=400, detail=str(e))
    finally:
        connection.close()