resmgmt.py

from fastapi import FastAPI, HTTPException, Depends, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.requests import Request
from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, Enum, Date, text, Float, DateTime
from sqlalchemy.orm import sessionmaker, declarative_base, relationship, Session
import enum
import datetime
from pydantic import BaseModel
from typing import List, Optional, Dict
from sqlalchemy import Enum as SqlEnum
from fastapi.exceptions import RequestValidationError


# ---------- FASTAPI APP ----------
app = FastAPI()

# ---------- ENUMS & SCHEMAS ----------

class TrainLocationUpdate(BaseModel):
    train_no: str
    latitude: float
    longitude: float
    timestamp: datetime.datetime
    
class TrainType(str, enum.Enum):
    EXPRESS = "Express"
    PASSENGER = "Passenger"
    FREIGHT = "Freight"

class TrainAllocationRequest(BaseModel):
    train_no: str
    engine_id: int
    pilot_id: Optional[int] = None
    carriage_ids: List[int]
    route: str
    start_date: datetime.date
    end_date: datetime.date
    train_type: TrainType = TrainType.EXPRESS


class EngineStatus(str, enum.Enum):
    AVAILABLE = "Available"
    IN_USE = "In Use"
    MAINTENANCE = "Maintenance"
    RESERVED = "Reserved"

class CarriageStatus(str, enum.Enum):
    AVAILABLE = "Available"
    IN_USE = "In Use"
    MAINTENANCE = "Maintenance"
    RESERVED = "Reserved"

class PilotStatus(str, enum.Enum):
    AVAILABLE = "Available"
    ASSIGNED = "Assigned"
    ON_LEAVE = "On Leave"
    TRAINING = "Training"
    RETIRED = "Retired"

class TrainStatus(str, enum.Enum):
    RUNNING = "Running"
    DELAYED = "Delayed"
    RESERVED = "Reserved"
    MAINTENANCE = "Maintenance"
    COMPLETED = "Completed"

# ---------- DATABASE CONFIG ----------
DATABASE_URL = "mysql+pymysql://root:mysql19t02r06@localhost:3306/resource_mgmt"

engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
Base = declarative_base()

# ---------- TABLE MODELS ----------

class TrainLocation(Base):
    __tablename__ = "train_location"
    id = Column(Integer, primary_key=True, index=True)
    train_no = Column(String(20), ForeignKey("train_allocation.train_no"))
    latitude = Column(Float)
    longitude = Column(Float)
    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
    train = relationship("TrainAllocation", back_populates="locations")

    
    
class Engine(Base):
    __tablename__ = "engine"
    engine_id = Column(Integer, primary_key=True, index=True)
    engine_type = Column(String(50))
    status = Column(Enum(EngineStatus, native_enum=False), default=EngineStatus.AVAILABLE)
    last_maintenance = Column(Date, nullable=True)
    next_maintenance = Column(Date, nullable=True)

class Carriage(Base):
    __tablename__ = "carriage"
    carriage_id = Column(Integer, primary_key=True, index=True)
    carriage_type = Column(String(50))
    station_id = Column(String(50))
    status = Column(Enum(CarriageStatus, native_enum=False), default=CarriageStatus.AVAILABLE)
    last_maintenance = Column(Date, nullable=True)
    next_maintenance = Column(Date, nullable=True)

class LocoPilot(Base):
    __tablename__ = "locopilot"
    pilot_id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100))
    experience_years = Column(Integer)
    status = Column(Enum(PilotStatus, native_enum=False, validate_strings=True), default=PilotStatus.AVAILABLE)
    last_trip_date = Column(Date, default=None)
    resting_period_days = Column(Integer, default=1)

class TrainAllocation(Base):
    __tablename__ = "train_allocation"
    train_no = Column(String(20), primary_key=True, index=True)
    engine_id = Column(Integer, ForeignKey("engine.engine_id"))
    pilot_id = Column(Integer, ForeignKey("locopilot.pilot_id"))
    route = Column(String(100))
    status = Column(Enum(TrainStatus, native_enum=False), default=TrainStatus.RESERVED)
    trip_start_date = Column(Date, default=None)
    trip_end_date = Column(Date, default=None)
    train_type = Column(SqlEnum(TrainType, native_enum=False), default=TrainType.EXPRESS)

    engine = relationship("Engine")
    pilot = relationship("LocoPilot")
    carriages = relationship("TrainCarriageMapping", back_populates="train")
    locations = relationship("TrainLocation", back_populates="train")

class TrainCarriageMapping(Base):
    __tablename__ = "train_carriage_mapping"
    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    train_no = Column(String(20), ForeignKey("train_allocation.train_no"))
    carriage_id = Column(Integer, ForeignKey("carriage.carriage_id"))

    train = relationship("TrainAllocation", back_populates="carriages")
    carriage = relationship("Carriage")

def calculate_pilot_score(pilot: LocoPilot, start_date: datetime.date) -> int:
    if pilot.status != PilotStatus.AVAILABLE:
        return -1

    if pilot.last_trip_date:
        if start_date < pilot.last_trip_date:
            return -1  # cannot start before last trip
        days_since_trip = (start_date - pilot.last_trip_date).days
        if days_since_trip < pilot.resting_period_days:
            return -1
        score = days_since_trip
    else:
        score = 5  # bonus for fresh availability

    score += pilot.experience_years * 2
    return score
   


# Dependency to get DB session
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

        
#---------web socket manager-----
class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, List[WebSocket]] = {}

    async def connect(self, train_no: str, websocket: WebSocket):
        await websocket.accept()
        if train_no not in self.active_connections:
            self.active_connections[train_no] = []
        self.active_connections[train_no].append(websocket)

    def disconnect(self, train_no: str, websocket: WebSocket):
        if train_no in self.active_connections:
            self.active_connections[train_no].remove(websocket)
            if not self.active_connections[train_no]:
                del self.active_connections[train_no]

    async def broadcast(self, message: Dict):
        for train_no, connections in self.active_connections.items():
            for connection in connections:
                try:
                    await connection.send_json(message)
                except WebSocketDisconnect:
                    self.disconnect(train_no, connection)


manager = ConnectionManager()

# ---------- API ENDPOINTS ----------
@app.get("/")
def root():
    return {"message": "Resource Management API (WCR Hackathon)"}

@app.get("/engines/")
def get_engines(db: Session = Depends(get_db)):
    return db.query(Engine).all()

@app.get("/carriages/")
def get_carriages(db: Session = Depends(get_db)):
    return db.query(Carriage).all()

@app.get("/pilots/")
def get_pilots(db: Session = Depends(get_db)):
    return db.query(LocoPilot).all()

# WebSocket endpoint for real-time tracking
@app.websocket("/ws/train-tracking/{train_no}")
async def websocket_endpoint(websocket: WebSocket, train_no: str):
    await manager.connect(train_no, websocket)
    try:
        while True:
            # We can receive messages here if the client needs to send data back
            # For now, we'll just keep the connection open.
            await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(train_no)

# API endpoint to receive location updates from a train's GPS/IoT device
@app.post("/api/update-location/")
async def update_train_location(location: TrainLocationUpdate, db: Session = Depends(get_db)):
    # 1. Save the location to the database
    new_location = TrainLocation(
        train_no=location.train_no,
        latitude=location.latitude,
        longitude=location.longitude,
        timestamp=location.timestamp
    )
    db.add(new_location)
    db.commit()
    db.refresh(new_location)
    
   # 2. Broadcast the location to all connected dashboard clients
    await manager.broadcast(location.dict())
    
    return {"message": "Location updated and broadcasted successfully"}

# Allocate train with conflict + rest + maintenance check
@app.post("/allocate-train/")
def allocate_train(request: TrainAllocationRequest, db: Session = Depends(get_db)):
    train_no = request.train_no
    engine_id = request.engine_id
    pilot_id = request.pilot_id
    carriage_ids = request.carriage_ids
    route = request.route
    start_date = request.start_date
    end_date = request.end_date

    # Check engine availability + maintenance
    engine = db.query(Engine).filter_by(engine_id=engine_id).first()
    if not engine or engine.status != EngineStatus.AVAILABLE:
        raise HTTPException(status_code=400, detail="Engine not available")
    if engine.next_maintenance and engine.next_maintenance < start_date:
        raise HTTPException(status_code=400, detail=f"Engine {engine.engine_id} requires maintenance before this trip")

    # Handle pilot selection
    if request.pilot_id:
        pilot = db.query(LocoPilot).filter_by(pilot_id=request.pilot_id).first()
        if not pilot:
            raise HTTPException(status_code=400, detail="Pilot not found")
        if pilot.status != PilotStatus.AVAILABLE:
            raise HTTPException(status_code=400, detail="Pilot not available")
        if pilot.last_trip_date:
            days_since_trip = (request.start_date - pilot.last_trip_date).days
            if days_since_trip < pilot.resting_period_days:
                raise HTTPException(
                    status_code=400,
                    detail=f"Pilot needs rest. Available after {pilot.last_trip_date + datetime.timedelta(days=pilot.resting_period_days)}"
                )
    else:
        # Auto-select best pilot
        candidates = db.query(LocoPilot).filter_by(status=PilotStatus.AVAILABLE).all()
        scored = [(p, calculate_pilot_score(p, request.start_date)) for p in candidates]
        scored = [p for p in scored if p[1] >= 0]
        if not scored:
            raise HTTPException(status_code=400, detail="No eligible pilots available")
        pilot = max(scored, key=lambda x: x[1])[0]

    # Check carriage availability + maintenance + overlap
    for cid in carriage_ids:
        carriage = db.query(Carriage).filter_by(carriage_id=cid).first()
        if not carriage or carriage.status != CarriageStatus.AVAILABLE:
            raise HTTPException(status_code=400, detail=f"Carriage {cid} not available")
        if carriage.next_maintenance and carriage.next_maintenance < start_date:
            raise HTTPException(status_code=400, detail=f"Carriage {cid} requires maintenance before this trip")
        overlap = db.query(TrainCarriageMapping).join(TrainAllocation).filter(
            TrainCarriageMapping.carriage_id == cid,
            TrainAllocation.trip_start_date <= end_date,
            TrainAllocation.trip_end_date >= start_date
        ).first()
        if overlap:
            raise HTTPException(status_code=400, detail=f"Carriage {cid} is already assigned in this period")

    # Create train allocation
    train = TrainAllocation(
        train_no=train_no,
        engine_id=engine_id,
        pilot_id=pilot.pilot_id,
        route=route,
        status=TrainStatus.RUNNING,
        trip_start_date=start_date,
        trip_end_date=end_date,
        train_type=request.train_type
    )
    db.add(train)
    db.flush()
    # Assign carriages
    for cid in carriage_ids:
        mapping = TrainCarriageMapping(train_no=train_no, carriage_id=cid)
        db.add(mapping)

    # Update resource statuses
    engine.status = EngineStatus.IN_USE
    pilot.status = PilotStatus.ASSIGNED
    pilot.last_trip_date = start_date
    for cid in carriage_ids:
        carriage = db.query(Carriage).filter_by(carriage_id=cid).first()
        carriage.status = CarriageStatus.IN_USE

    db.commit()
    return {
        "message": "Train allocated successfully",
        "train_no": train_no,
        "train_type": request.train_type.value,
        "engine_id": engine_id,
        "pilot_id": pilot.pilot_id,
        "pilot_name": pilot.name,
        "carriage_ids": carriage_ids,
        "route": route,
        "start_date": str(start_date),
        "end_date": str(end_date)
    }

# Reset resources after train completes journey
@app.post("/reset-train/{train_no}")
def reset_train(train_no: str, db: Session = Depends(get_db)):
    train = db.query(TrainAllocation).filter_by(train_no=train_no).first()
    if not train:
        raise HTTPException(404, "Train not found")

    train.status = TrainStatus.COMPLETED

    # Reset engine
    if train.engine:
        train.engine.status = EngineStatus.AVAILABLE
        # update maintenance
        train.engine.last_maintenance = train.trip_end_date
        train.engine.next_maintenance = train.trip_end_date + datetime.timedelta(days=30)  # example rule: every 30 days

    # Reset pilot (respecting rest days)
    if train.pilot:
        train.pilot.last_trip_date = train.trip_end_date
        train.pilot.status = PilotStatus.AVAILABLE
        

    # Reset carriages
    for mapping in train.carriages:
        mapping.carriage.status = CarriageStatus.AVAILABLE
        # update maintenance
        mapping.carriage.last_maintenance = train.trip_end_date
        mapping.carriage.next_maintenance = train.trip_end_date + datetime.timedelta(days=45)  # example: every 45 days
    db.commit()
    return {"message": f"Train {train_no} reset successfully"}



@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    if isinstance(exc, HTTPException):
        return JSONResponse(
            status_code=exc.status_code,
            content={"message": exc.detail},
        )

    return JSONResponse(
        status_code=500,
        content={"message": f"Unexpected error: {str(exc)}"}
    )




# ---------- HEALTH CHECK ----------
@app.get("/health")
def health_check(db: Session = Depends(get_db)):
    try:
        db.execute(text("SELECT 1"))
        return {"status": "ok"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))