terra-view/backend/routers/roster_edit.py

from fastapi import APIRouter, Depends, HTTPException, Form, UploadFile, File
from sqlalchemy.orm import Session
from datetime import datetime
import csv
import io

from backend.database import get_db
from backend.models import RosterUnit

router = APIRouter(prefix="/api/roster", tags=["roster-edit"])


def get_or_create_roster_unit(db: Session, unit_id: str):
    unit = db.query(RosterUnit).filter(RosterUnit.id == unit_id).first()
    if not unit:
        unit = RosterUnit(id=unit_id)
        db.add(unit)
        db.commit()
        db.refresh(unit)
    return unit


@router.post("/add")
def add_roster_unit(
    id: str = Form(...),
    unit_type: str = Form("series3"),
    deployed: bool = Form(True),
    note: str = Form(""),
    db: Session = Depends(get_db)
):
    if db.query(RosterUnit).filter(RosterUnit.id == id).first():
        raise HTTPException(status_code=400, detail="Unit already exists")

    unit = RosterUnit(
        id=id,
        unit_type=unit_type,
        deployed=deployed,
        note=note,
        last_updated=datetime.utcnow(),
    )
    db.add(unit)
    db.commit()
    return {"message": "Unit added", "id": id}


@router.post("/set-deployed/{unit_id}")
def set_deployed(unit_id: str, deployed: bool = Form(...), db: Session = Depends(get_db)):
    unit = get_or_create_roster_unit(db, unit_id)
    unit.deployed = deployed
    unit.last_updated = datetime.utcnow()
    db.commit()
    return {"message": "Updated", "id": unit_id, "deployed": deployed}


@router.post("/set-retired/{unit_id}")
def set_retired(unit_id: str, retired: bool = Form(...), db: Session = Depends(get_db)):
    unit = get_or_create_roster_unit(db, unit_id)
    unit.retired = retired
    unit.last_updated = datetime.utcnow()
    db.commit()
    return {"message": "Updated", "id": unit_id, "retired": retired}


@router.post("/set-note/{unit_id}")
def set_note(unit_id: str, note: str = Form(""), db: Session = Depends(get_db)):
    unit = get_or_create_roster_unit(db, unit_id)
    unit.note = note
    unit.last_updated = datetime.utcnow()
    db.commit()
    return {"message": "Updated", "id": unit_id, "note": note}


@router.post("/import-csv")
async def import_csv(
    file: UploadFile = File(...),
    update_existing: bool = Form(True),
    db: Session = Depends(get_db)
):
    """
    Import roster units from CSV file.

    Expected CSV columns (unit_id is required, others are optional):
    - unit_id: Unique identifier for the unit
    - unit_type: Type of unit (default: "series3")
    - deployed: Boolean for deployment status (default: False)
    - retired: Boolean for retirement status (default: False)
    - note: Notes about the unit
    - project_id: Project identifier
    - location: Location description

    Args:
        file: CSV file upload
        update_existing: If True, update existing units; if False, skip them
    """

    if not file.filename.endswith('.csv'):
        raise HTTPException(status_code=400, detail="File must be a CSV")

    # Read file content
    contents = await file.read()
    csv_text = contents.decode('utf-8')
    csv_reader = csv.DictReader(io.StringIO(csv_text))

    results = {
        "added": [],
        "updated": [],
        "skipped": [],
        "errors": []
    }

    for row_num, row in enumerate(csv_reader, start=2):  # Start at 2 to account for header
        try:
            # Validate required field
            unit_id = row.get('unit_id', '').strip()
            if not unit_id:
                results["errors"].append({
                    "row": row_num,
                    "error": "Missing required field: unit_id"
                })
                continue

            # Check if unit exists
            existing_unit = db.query(RosterUnit).filter(RosterUnit.id == unit_id).first()

            if existing_unit:
                if not update_existing:
                    results["skipped"].append(unit_id)
                    continue

                # Update existing unit
                existing_unit.unit_type = row.get('unit_type', existing_unit.unit_type or 'series3')
                existing_unit.deployed = row.get('deployed', '').lower() in ('true', '1', 'yes') if row.get('deployed') else existing_unit.deployed
                existing_unit.retired = row.get('retired', '').lower() in ('true', '1', 'yes') if row.get('retired') else existing_unit.retired
                existing_unit.note = row.get('note', existing_unit.note or '')
                existing_unit.project_id = row.get('project_id', existing_unit.project_id)
                existing_unit.location = row.get('location', existing_unit.location)
                existing_unit.last_updated = datetime.utcnow()

                results["updated"].append(unit_id)
            else:
                # Create new unit
                new_unit = RosterUnit(
                    id=unit_id,
                    unit_type=row.get('unit_type', 'series3'),
                    deployed=row.get('deployed', '').lower() in ('true', '1', 'yes'),
                    retired=row.get('retired', '').lower() in ('true', '1', 'yes'),
                    note=row.get('note', ''),
                    project_id=row.get('project_id'),
                    location=row.get('location'),
                    last_updated=datetime.utcnow()
                )
                db.add(new_unit)
                results["added"].append(unit_id)

        except Exception as e:
            results["errors"].append({
                "row": row_num,
                "unit_id": row.get('unit_id', 'unknown'),
                "error": str(e)
            })

    # Commit all changes
    try:
        db.commit()
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")

    return {
        "message": "CSV import completed",
        "summary": {
            "added": len(results["added"]),
            "updated": len(results["updated"]),
            "skipped": len(results["skipped"]),
            "errors": len(results["errors"])
        },
        "details": results
    }