TJWaterServerBinary/app/api/v1/endpoints/model_import.py

import json
from pathlib import Path
from tempfile import NamedTemporaryFile
from uuid import UUID, uuid4

from fastapi import (
    APIRouter,
    Depends,
    File,
    Form,
    HTTPException,
    Path as ApiPath,
    Request,
    UploadFile,
    status,
)
from sqlalchemy.exc import IntegrityError
from starlette.concurrency import run_in_threadpool

from app.auth.metadata_dependencies import (
    get_current_metadata_admin,
    get_metadata_repository,
)
from app.auth.project_dependencies import (
    ProjectContext,
    resolve_project_business_routing,
)
from app.core.audit import AuditAction, log_audit_event
from app.core.encryption import is_database_encryption_configured
from app.domain.schemas.admin_metadata import (
    AdminProjectResponse,
    ProjectProvisionResponse,
)
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
from app.infra.db.project_routing import activate_project_routing
from app.native.wndb.core.database import MaterializedViewRefreshAfterCommitError
from app.services.network_import import network_update
from app.services.project_provisioning import (
    ProjectProvisioningError,
    ProvisionedProjectInfrastructure,
    provision_project_infrastructure,
    validate_project_code,
)
from app.services.tjnetwork import run_inp

router = APIRouter()

MAX_INP_FILE_BYTES = 50 * 1024 * 1024
INP_SECTIONS = ("[TITLE]", "[JUNCTIONS]", "[RESERVOIRS]", "[TANKS]", "[PIPES]")


async def _get_active_project(project_id: UUID, metadata_repo: MetadataRepository):
    project = await metadata_repo.get_project_by_id(project_id)
    if project is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Project not found",
        )
    if project.status != "active":
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Project is not active",
        )
    return project


def _validate_inp_bytes(content: bytes, filename: str) -> str:
    if Path(filename).suffix.lower() != ".inp":
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Only .inp model files are accepted",
        )
    if not content:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="INP file is empty",
        )
    if len(content) > MAX_INP_FILE_BYTES:
        raise HTTPException(
            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
            detail="INP file exceeds the 50 MiB limit",
        )
    for encoding in ("utf-8-sig", "gb18030"):
        try:
            text = content.decode(encoding)
            break
        except UnicodeDecodeError:
            continue
    else:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="INP file encoding is not supported",
        )
    upper_text = text.upper()
    if not any(section in upper_text for section in INP_SECTIONS):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid INP file structure",
        )
    return text


async def _read_upload(file: UploadFile) -> tuple[bytes, str]:
    filename = Path(file.filename or "").name
    content = await file.read(MAX_INP_FILE_BYTES + 1)
    normalized = _validate_inp_bytes(content, filename).encode("utf-8")
    return normalized, filename


async def _audit_model_change(
    *,
    request: Request,
    current_user,
    metadata_repo: MetadataRepository,
    project_id: UUID,
    action: str,
) -> None:
    await log_audit_event(
        action=AuditAction.UPDATE,
        user_id=current_user.id,
        project_id=project_id,
        resource_type="hydraulic_model",
        resource_id=action,
        request_data={"operation": action},
        ip_address=request.client.host if request.client else None,
        request_method=request.method,
        request_path=request.url.path,
        response_status=status.HTTP_200_OK,
        session=metadata_repo.session,
    )


def _run_uploaded_inp_sync(content: bytes) -> str:
    target_dir = Path("inp")
    target_dir.mkdir(parents=True, exist_ok=True)
    model_name = f"admin_model_{uuid4().hex}"
    target_path = target_dir / f"{model_name}.inp"
    target_path.write_bytes(content)
    return run_inp(model_name)


async def _run_uploaded_inp(content: bytes) -> str:
    return await run_in_threadpool(_run_uploaded_inp_sync, content)


def _update_from_inp_sync(content: bytes, project_code: str) -> None:
    temp_path: Path | None = None
    try:
        with NamedTemporaryFile(suffix=".inp", delete=False) as temp_file:
            temp_file.write(content)
            temp_path = Path(temp_file.name)
        network_update(str(temp_path), project_code)
    finally:
        if temp_path is not None:
            temp_path.unlink(missing_ok=True)


async def _update_from_inp(content: bytes, project_code: str) -> None:
    await run_in_threadpool(_update_from_inp_sync, content, project_code)


async def _apply_model_update(content: bytes, project_code: str) -> None:
    try:
        await _update_from_inp(content, project_code)
    except MaterializedViewRefreshAfterCommitError:
        raise
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"数据库操作失败: {exc}",
        ) from exc


def _provision_from_inp_sync(
    content: bytes,
    *,
    code: str,
    workspace: str,
) -> ProvisionedProjectInfrastructure:
    temp_path: Path | None = None
    try:
        with NamedTemporaryFile(suffix=".inp", delete=False) as temp_file:
            temp_file.write(content)
            temp_path = Path(temp_file.name)
        return provision_project_infrastructure(
            code=code,
            workspace=workspace,
            inp_path=temp_path,
        )
    finally:
        if temp_path is not None:
            temp_path.unlink(missing_ok=True)


@router.post(
    "/admin/project-provisions",
    response_model=ProjectProvisionResponse,
    status_code=status.HTTP_201_CREATED,
    summary="创建完整供水项目",
)
async def provision_project(
    request: Request,
    name: str = Form(..., min_length=1, max_length=100),
    code: str = Form(..., min_length=1, max_length=50),
    description: str | None = Form(default=None),
    gs_workspace: str | None = Form(default=None, max_length=100),
    map_zoom: int = Form(default=14, ge=1, le=22),
    file: UploadFile = File(..., description="EPANET INP 模型文件"),
    current_user=Depends(get_current_metadata_admin),
    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> ProjectProvisionResponse:
    try:
        normalized_code = validate_project_code(code)
    except ValueError as exc:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=str(exc),
        ) from exc
    workspace = gs_workspace or normalized_code
    if await metadata_repo.get_project_by_code(normalized_code) is not None:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="Project code already exists",
        )
    if not is_database_encryption_configured():
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="DATABASE_ENCRYPTION_KEY is not configured",
        )

    content, filename = await _read_upload(file)
    validation_result = await _run_uploaded_inp(content)
    try:
        validation_payload = json.loads(validation_result)
    except (TypeError, json.JSONDecodeError) as exc:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="EPANET validation returned an invalid response",
        ) from exc
    if validation_payload.get("simulation_result") != "successful":
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="EPANET model validation failed",
        )
    try:
        infrastructure = await run_in_threadpool(
            _provision_from_inp_sync,
            content,
            code=normalized_code,
            workspace=workspace,
        )
    except ProjectProvisioningError as exc:
        if isinstance(exc.cause, ValueError):
            response_status = status.HTTP_409_CONFLICT
        elif exc.stage == "preflight":
            response_status = status.HTTP_503_SERVICE_UNAVAILABLE
        else:
            response_status = status.HTTP_500_INTERNAL_SERVER_ERROR
        raise HTTPException(
            status_code=response_status,
            detail={
                "stage": exc.stage,
                "message": str(exc.cause),
                "cleanup_errors": exc.cleanup_errors,
            },
        ) from exc

    map_extent = {"bbox": list(infrastructure.map_bbox), "zoom": map_zoom}
    try:
        project = await metadata_repo.create_provisioned_project(
            name=name,
            code=normalized_code,
            description=description,
            gs_workspace=workspace,
            map_extent=map_extent,
            creator_user_id=current_user.id,
            business_dsn=infrastructure.business_dsn,
            timescale_dsn=infrastructure.timescale_dsn,
            pool_min_size=1,
            pool_max_size=4,
        )
    except Exception as exc:
        await metadata_repo.session.rollback()
        cleanup_errors = await run_in_threadpool(infrastructure.cleanup)
        if isinstance(exc, IntegrityError):
            response_status = status.HTTP_409_CONFLICT
            detail = "Project code or workspace conflicts with an existing project"
        else:
            response_status = status.HTTP_503_SERVICE_UNAVAILABLE
            detail = f"Metadata database error: {exc}"
        if cleanup_errors:
            detail = f"{detail}; cleanup failures: {', '.join(cleanup_errors)}"
        raise HTTPException(status_code=response_status, detail=detail) from exc

    await log_audit_event(
        action=AuditAction.CREATE,
        user_id=current_user.id,
        project_id=project.id,
        resource_type="project_provision",
        resource_id=str(project.id),
        request_data={
            "name": name,
            "code": normalized_code,
            "filename": filename,
            "gs_workspace": workspace,
            "layers": list(infrastructure.layers),
        },
        ip_address=request.client.host if request.client else None,
        request_method=request.method,
        request_path=request.url.path,
        response_status=status.HTTP_201_CREATED,
        session=metadata_repo.session,
    )
    return ProjectProvisionResponse(
        project=AdminProjectResponse(
            project_id=project.id,
            name=project.name,
            code=project.code,
            description=project.description,
            gs_workspace=project.gs_workspace,
            map_extent=project.map_extent,
            status=project.status,
            created_at=project.created_at,
            updated_at=project.updated_at,
        ),
        business_database=normalized_code,
        model_template_database=infrastructure.model_template,
        timescale_database=normalized_code,
        geoserver_workspace=workspace,
        geoserver_layers=list(infrastructure.layers),
    )


@router.post(
    "/admin/projects/{project_id}/model-imports",
    summary="导入桌面端水力模型",
)
async def import_project_model(
    request: Request,
    project_id: UUID = ApiPath(...),
    file: UploadFile = File(..., description="桌面端导出的 INP 模型文件"),
    current_user=Depends(get_current_metadata_admin),
    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> dict:
    project = await _get_active_project(project_id, metadata_repo)
    content, filename = await _read_upload(file)
    result = await _run_uploaded_inp(content)
    await _audit_model_change(
        request=request,
        current_user=current_user,
        metadata_repo=metadata_repo,
        project_id=project.id,
        action="import",
    )
    return {"project_id": str(project.id), "filename": filename, "result": result}


@router.patch(
    "/admin/projects/{project_id}/model-imports",
    summary="更新桌面端水力模型",
)
async def update_project_model(
    request: Request,
    project_id: UUID = ApiPath(...),
    file: UploadFile = File(..., description="桌面端导出的 INP 模型文件"),
    current_user=Depends(get_current_metadata_admin),
    metadata_repo: MetadataRepository = Depends(get_metadata_repository),
) -> dict:
    project = await _get_active_project(project_id, metadata_repo)
    content, filename = await _read_upload(file)
    routing = await resolve_project_business_routing(
        ProjectContext(
            project_id=project.id,
            project_code=project.code,
            user_id=current_user.id,
            project_role="owner",
            system_role=current_user.role,
            is_superuser=current_user.is_superuser,
        ),
        metadata_repo,
    )
    with activate_project_routing(routing):
        await _apply_model_update(content, project.code)
    await _audit_model_change(
        request=request,
        current_user=current_user,
        metadata_repo=metadata_repo,
        project_id=project.id,
        action="update",
    )
    return {"project_id": str(project.id), "filename": filename, "updated": True}