cosmo/backend/app/api/rocket.py

"""Rocket simulator configuration APIs."""
from datetime import datetime
from typing import Annotated

from fastapi import APIRouter, Depends, HTTPException, Query, status
from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator
from sqlalchemy import func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import get_db
from app.models.db import RocketConfig, User
from app.services.auth_deps import require_admin


router = APIRouter(prefix="/rockets", tags=["rockets"])
RocketCode = Annotated[str, StringConstraints(pattern=r"^[a-z0-9][a-z0-9-]*$")]


class RocketStage(BaseModel):
    name: str = Field(min_length=1, max_length=100)
    dry_mass_kg: float = Field(gt=0)
    fuel_mass_kg: float = Field(gt=0)
    max_thrust_n: float = Field(gt=0)
    specific_impulse_s: float = Field(gt=0)
    engine_count: int = Field(gt=0)
    length_ratio: float = Field(gt=0, le=1)


class RocketBase(BaseModel):
    code: RocketCode
    name: str = Field(min_length=1, max_length=100)
    name_zh: str | None = Field(default=None, max_length=100)
    manufacturer: str | None = Field(default=None, max_length=100)
    country: str | None = Field(default=None, max_length=100)
    launch_site_name: str = Field(default="Equatorial Launch Site", min_length=1, max_length=120)
    launch_latitude_deg: float = Field(default=0, ge=-90, le=90)
    launch_longitude_deg: float = Field(default=0, ge=-180, le=180)
    description: str | None = None
    color: str = Field(default="#f8fafc", pattern=r"^#[0-9a-fA-F]{6}$")
    height_m: float = Field(gt=0)
    diameter_m: float = Field(gt=0)
    payload_mass_kg: float = Field(ge=0, default=0)
    drag_coefficient: float = Field(gt=0, le=2, default=0.4)
    reference_area_m2: float = Field(gt=0)
    target_orbit_km: float = Field(gt=0, default=200)
    target_velocity_mps: float = Field(gt=0, default=7800)
    separation_delay_seconds: float = Field(ge=0, le=30, default=2)
    second_stage_ignition_delay_seconds: float = Field(ge=0, le=30, default=1)
    stage_1: RocketStage
    stage_2: RocketStage
    is_active: bool = True
    sort_order: int = 0

    @field_validator("code", mode="before")
    @classmethod
    def normalize_code(cls, value: str) -> str:
        return value.strip().lower()


class RocketCreate(RocketBase):
    pass


class RocketUpdate(BaseModel):
    code: RocketCode | None = None
    name: str | None = Field(default=None, min_length=1, max_length=100)
    name_zh: str | None = Field(default=None, max_length=100)
    manufacturer: str | None = Field(default=None, max_length=100)
    country: str | None = Field(default=None, max_length=100)
    launch_site_name: str | None = Field(default=None, min_length=1, max_length=120)
    launch_latitude_deg: float | None = Field(default=None, ge=-90, le=90)
    launch_longitude_deg: float | None = Field(default=None, ge=-180, le=180)
    description: str | None = None
    color: str | None = Field(default=None, pattern=r"^#[0-9a-fA-F]{6}$")
    height_m: float | None = Field(default=None, gt=0)
    diameter_m: float | None = Field(default=None, gt=0)
    payload_mass_kg: float | None = Field(default=None, ge=0)
    drag_coefficient: float | None = Field(default=None, gt=0, le=2)
    reference_area_m2: float | None = Field(default=None, gt=0)
    target_orbit_km: float | None = Field(default=None, gt=0)
    target_velocity_mps: float | None = Field(default=None, gt=0)
    separation_delay_seconds: float | None = Field(default=None, ge=0, le=30)
    second_stage_ignition_delay_seconds: float | None = Field(default=None, ge=0, le=30)
    stage_1: RocketStage | None = None
    stage_2: RocketStage | None = None
    is_active: bool | None = None
    sort_order: int | None = None

    @field_validator("code", mode="before")
    @classmethod
    def normalize_code(cls, value: str | None) -> str | None:
        return value.strip().lower() if value else value


class RocketResponse(RocketBase):
    id: int
    created_at: datetime | None = None
    updated_at: datetime | None = None

    model_config = ConfigDict(from_attributes=True)


def rocket_values(data: RocketCreate | RocketUpdate, *, exclude_unset: bool = False) -> dict:
    values = data.model_dump(exclude_unset=exclude_unset)
    for key in ("stage_1", "stage_2"):
        if key in values and values[key] is not None:
            values[key] = dict(values[key])
    return values


async def ensure_unique_code(db: AsyncSession, code: str, exclude_id: int | None = None) -> None:
    query = select(RocketConfig.id).where(RocketConfig.code == code)
    if exclude_id is not None:
        query = query.where(RocketConfig.id != exclude_id)
    if (await db.execute(query)).scalar_one_or_none() is not None:
        raise HTTPException(status_code=400, detail="Rocket code already exists")


@router.get("", response_model=list[RocketResponse])
async def list_active_rockets(db: AsyncSession = Depends(get_db)):
    result = await db.execute(
        select(RocketConfig)
        .where(RocketConfig.is_active.is_(True))
        .order_by(RocketConfig.sort_order, RocketConfig.id)
    )
    return result.scalars().all()


@router.get("/admin")
async def list_rockets_for_admin(
    skip: int = Query(0, ge=0),
    limit: int = Query(20, ge=1, le=100),
    search: str | None = None,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_admin),
):
    filters = []
    if search:
        keyword = f"%{search.strip()}%"
        filters.append(or_(
            RocketConfig.code.ilike(keyword),
            RocketConfig.name.ilike(keyword),
            RocketConfig.name_zh.ilike(keyword),
        ))

    total = await db.scalar(select(func.count(RocketConfig.id)).where(*filters))
    result = await db.execute(
        select(RocketConfig)
        .where(*filters)
        .order_by(RocketConfig.sort_order, RocketConfig.id)
        .offset(skip)
        .limit(limit)
    )
    return {
        "rockets": [RocketResponse.model_validate(item) for item in result.scalars().all()],
        "total": total or 0,
    }


@router.get("/{code}", response_model=RocketResponse)
async def get_active_rocket(code: str, db: AsyncSession = Depends(get_db)):
    result = await db.execute(
        select(RocketConfig).where(
            RocketConfig.code == code.lower(),
            RocketConfig.is_active.is_(True),
        )
    )
    rocket = result.scalar_one_or_none()
    if not rocket:
        raise HTTPException(status_code=404, detail="Rocket not found")
    return rocket


@router.post("/admin", response_model=RocketResponse, status_code=status.HTTP_201_CREATED)
async def create_rocket(
    data: RocketCreate,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_admin),
):
    await ensure_unique_code(db, data.code)
    rocket = RocketConfig(**rocket_values(data))
    db.add(rocket)
    await db.commit()
    await db.refresh(rocket)
    return rocket


@router.put("/admin/{rocket_id}", response_model=RocketResponse)
async def update_rocket(
    rocket_id: int,
    data: RocketUpdate,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_admin),
):
    rocket = await db.get(RocketConfig, rocket_id)
    if not rocket:
        raise HTTPException(status_code=404, detail="Rocket not found")

    values = rocket_values(data, exclude_unset=True)
    if not values:
        raise HTTPException(status_code=400, detail="No fields to update")
    if "code" in values:
        await ensure_unique_code(db, values["code"], rocket_id)

    for key, value in values.items():
        setattr(rocket, key, value)
    await db.commit()
    await db.refresh(rocket)
    return rocket


@router.delete("/admin/{rocket_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_rocket(
    rocket_id: int,
    db: AsyncSession = Depends(get_db),
    _: User = Depends(require_admin),
):
    rocket = await db.get(RocketConfig, rocket_id)
    if not rocket:
        raise HTTPException(status_code=404, detail="Rocket not found")
    await db.delete(rocket)
    await db.commit()