DotsToSirface/backend/main.py

from __future__ import annotations

import json
import os
import struct
import subprocess
import tempfile
import uuid
from pathlib import Path
from typing import Any

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from builtin_presets import BUILTIN_PRESETS, get_builtin_preset
from demo_generator import DEMO_SURFACE_TYPES, demo_payload
from pipeline_insights import compute_insights
from dataset_generator import generate_dataset, load_object_points_from_obj_text, load_scene_preview
from scene_generator import (
    catalog_payload as generator_catalog_payload,
    export_npy_float64,
    export_obj,
    export_ply,
    export_xyz,
    generate_layer,
    layers_to_pointnet_rows,
    merge_layers_world,
    parse_obj_points,
    points_to_pointnet_rows,
    resolve_intersections,
)
from stage_meta import STAGE_META, STAGE_META_BY_ID, catalog_payload, defaults_for_stage

APP_ROOT = Path(__file__).resolve().parent.parent
WEB_DIST = APP_ROOT / "frontend" / "web" / "dist"
PRESETS_DIRS = [APP_ROOT / "presets", APP_ROOT / "docker"]
USER_PRESETS_DIR = Path(os.environ.get("USER_PRESETS_DIR", "/app/data/user-presets"))
DEFAULT_PIPELINE_CONFIG = Path(
    os.environ.get("PIPELINE_CONFIG", APP_ROOT / "docker" / "default_pipeline.json")
)
DOTSTOSURFACE_BIN = Path(os.environ.get("DOTSTOSURFACE_BIN", "/usr/local/bin/DotsToSurface"))

app = FastAPI(title="DotsToSurface Web API", version="2.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

USER_PRESETS_DIR.mkdir(parents=True, exist_ok=True)


class PipelineConfigBody(BaseModel):
    profile: str = "desktop_debug"
    preprocessPlugins: list[str] = []
    reconstructionPlugin: str = "surface_fallback"
    stageDefaults: dict[str, str] = {}


class ValidateConfigBody(BaseModel):
    config: PipelineConfigBody


class WizardBody(BaseModel):
    wizardProfile: str = "general"
    wizardGoal: str = "balanced"


class UserPresetBody(BaseModel):
    title: str
    idValue: str | None = None
    stages: list[dict[str, Any]]


class GeneratorLayerBody(BaseModel):
    kind: str
    type: str
    params: dict[str, Any] | None = None


class GeneratorLayerItem(BaseModel):
    id: str | None = None
    kind: str
    type: str
    name: str | None = None
    params: dict[str, Any] | None = None
    transform: dict[str, Any] | None = None
    points: list[list[float]] | None = None
    label: str | None = None
    color: str | None = None


class GeneratorResolveBody(BaseModel):
    layers: list[GeneratorLayerItem]
    eps: float = 0.01
    clipSurfaceInsideObjects: bool = True
    clipObjectsVsObjects: bool = True


class GeneratorExportBody(BaseModel):
    points: list[list[float]] | None = None
    layers: list[GeneratorLayerItem] | None = None
    format: str = "xyz"
    filename: str | None = None
    # For points-only .npy export when layers are not provided (1=pipe, 0=other).
    classLabel: float | None = None


class DatasetGenerateBody(BaseModel):
    count: int = 5
    seed: int = 42
    outputDir: str = "sonar_dataset"


class DatasetPreviewBody(BaseModel):
    stem: str
    outputDir: str = "sonar_dataset"
    maxPoints: int = 25000


def preset_to_pipeline_config(preset: dict[str, Any]) -> dict[str, Any]:
    if "preprocessPlugins" in preset and "reconstructionPlugin" in preset:
        return preset
    if "config" in preset and isinstance(preset["config"], dict):
        return preset["config"]

    preprocess_plugins: list[str] = []
    reconstruction_plugin = "surface_fallback"
    stage_defaults: dict[str, str] = {}

    for stage in preset.get("stages", []):
        if not stage.get("enabled", True):
            continue
        stage_id = str(stage.get("id", "")).strip()
        if not stage_id:
            continue
        family = str(stage.get("family", "")).strip()
        if family == "preprocess":
            preprocess_plugins.append(stage_id)
        elif family == "reconstruction":
            reconstruction_plugin = stage_id
        defaults = str(stage.get("defaults", "")).strip()
        if defaults:
            stage_defaults[stage_id] = defaults

    return {
        "profile": preset.get("profile", "desktop_debug"),
        "preprocessPlugins": preprocess_plugins,
        "reconstructionPlugin": reconstruction_plugin,
        "stageDefaults": stage_defaults,
    }


def config_to_stage_cards(config: dict[str, Any]) -> list[dict[str, Any]]:
    cards: list[dict[str, Any]] = []
    stage_defaults = config.get("stageDefaults", {})
    for stage_id in config.get("preprocessPlugins", []):
        meta = STAGE_META_BY_ID.get(stage_id, {})
        cards.append(
            {
                "id": stage_id,
                "title": meta.get("title", stage_id),
                "category": meta.get("category", "Custom"),
                "family": "preprocess",
                "hint": meta.get("hint", ""),
                "defaults": stage_defaults.get(stage_id, meta.get("defaults", "")),
                "enabled": True,
            }
        )
    recon_id = config.get("reconstructionPlugin", "surface_fallback")
    recon_meta = STAGE_META_BY_ID.get(recon_id, {})
    cards.append(
        {
            "id": recon_id,
            "title": recon_meta.get("title", recon_id),
            "category": recon_meta.get("category", "Реконструкция"),
            "family": "reconstruction",
            "hint": recon_meta.get("hint", ""),
            "defaults": stage_defaults.get(recon_id, recon_meta.get("defaults", "")),
            "enabled": True,
        }
    )
    return cards


def list_preset_files() -> list[Path]:
    files: list[Path] = []
    seen: set[str] = set()
    for directory in PRESETS_DIRS:
        if not directory.is_dir():
            continue
        for path in sorted(directory.glob("*.json")):
            if path.name in seen:
                continue
            seen.add(path.name)
            files.append(path)
    return files


def user_presets_path() -> Path:
    return USER_PRESETS_DIR / "pipeline_presets.json"


def load_user_presets() -> list[dict[str, Any]]:
    path = user_presets_path()
    if not path.is_file():
        return []
    with path.open("r", encoding="utf-8") as handle:
        data = json.load(handle)
    return data if isinstance(data, list) else []


def save_user_presets(presets: list[dict[str, Any]]) -> None:
    with user_presets_path().open("w", encoding="utf-8") as handle:
        json.dump(presets, handle, indent=2, ensure_ascii=False)


def write_binary_geometry(work_dir: Path, result: dict[str, Any]) -> str | None:
    points = result.get("points") or []
    triangles = result.get("triangleIndices") or []
    if not points:
        return None
    bin_path = work_dir / "geometry.bin"
    with bin_path.open("wb") as handle:
        handle.write(struct.pack("<I", len(points)))
        for point in points:
            handle.write(struct.pack("<3f", float(point[0]), float(point[1]), float(point[2])))
        handle.write(struct.pack("<I", len(triangles)))
        for tri in triangles:
            handle.write(struct.pack("<3i", int(tri[0]), int(tri[1]), int(tri[2])))
    return str(bin_path)


def run_pipeline_command(input_path: Path, config: dict[str, Any], output_path: Path) -> subprocess.CompletedProcess[str]:
    config_path = output_path.parent / "pipeline_config.json"
    config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
    command = [
        str(DOTSTOSURFACE_BIN),
        "--cli",
        "--input",
        str(input_path),
        "--config-json",
        str(config_path),
        "--output-json",
        str(output_path),
    ]
    return subprocess.run(command, capture_output=True, text=True, check=False)


def enrich_result(result: dict[str, Any], config: dict[str, Any], stdout: str, work_id: str) -> dict[str, Any]:
    insights = compute_insights(
        config.get("preprocessPlugins", []),
        config.get("reconstructionPlugin", "surface_fallback"),
    )
    result.update(insights)
    result["stdout"] = stdout
    result["workId"] = work_id
    result["stageCards"] = config_to_stage_cards(config)
    result.setdefault(
        "metrics",
        {
            "inputPoints": result.get("inputPoints", 0),
            "afterPreprocess": result.get("afterPreprocess", 0),
            "triangles": result.get("triangles", 0),
            "reconstructMs": result.get("reconstructionMs", 0),
            "clusters": result.get("clusters", 0),
            "removedPoints": result.get("removedPoints", 0),
        },
    )
    return result


@app.get("/api/health")
def health() -> dict[str, str]:
    return {
        "status": "ok",
        "binary": str(DOTSTOSURFACE_BIN),
        "binaryExists": str(DOTSTOSURFACE_BIN.is_file()),
        "webDist": str(WEB_DIST.is_dir()),
    }


@app.get("/api/catalog")
def catalog() -> dict[str, Any]:
    return catalog_payload()


@app.get("/api/builtin-presets")
def builtin_presets() -> list[dict[str, Any]]:
    return BUILTIN_PRESETS


@app.get("/api/presets")
def presets() -> list[dict[str, Any]]:
    items = list(BUILTIN_PRESETS)
    for path in list_preset_files():
        with path.open("r", encoding="utf-8") as handle:
            data = json.load(handle)
        items.append(
            {
                "id": path.stem,
                "filename": path.name,
                "title": data.get("title", path.stem),
                "idValue": data.get("idValue", path.stem),
                "config": preset_to_pipeline_config(data),
                "stages": data.get("stages"),
            }
        )
    for preset in load_user_presets():
        items.append(
            {
                "id": preset.get("idValue", ""),
                "title": preset.get("title", "User preset"),
                "idValue": preset.get("idValue", ""),
                "config": preset_to_pipeline_config(preset),
                "stages": preset.get("stages"),
                "user": True,
            }
        )
    return items


@app.get("/api/default-config")
def default_config() -> dict[str, Any]:
    if not DEFAULT_PIPELINE_CONFIG.is_file():
        raise HTTPException(status_code=500, detail="Default pipeline config is missing.")
    with DEFAULT_PIPELINE_CONFIG.open("r", encoding="utf-8") as handle:
        return json.load(handle)


@app.get("/api/stage-defaults/{stage_id}")
def stage_defaults(stage_id: str) -> dict[str, str]:
    return {"stageId": stage_id, "defaults": defaults_for_stage(stage_id)}


@app.post("/api/validate-config")
def validate_config(body: ValidateConfigBody) -> dict[str, Any]:
    config = body.config.model_dump()
    insights = compute_insights(config.get("preprocessPlugins", []), config.get("reconstructionPlugin", ""))
    return {
        "config": config,
        "stageCards": config_to_stage_cards(config),
        **insights,
    }


@app.post("/api/wizard")
def wizard_suggestion(body: WizardBody) -> dict[str, Any]:
    preset_id = "Synthetic_clean"
    if body.wizardProfile == "urban_scan":
        preset_id = "LiDAR_scan"
    elif body.wizardProfile == "indoor_object":
        preset_id = "RGBD_camera"
    if body.wizardGoal == "speed":
        if preset_id == "LiDAR_scan":
            preset_id = "Synthetic_clean"
        elif preset_id == "RGBD_camera":
            preset_id = "Fast"
    elif body.wizardGoal == "quality":
        if preset_id == "Synthetic_clean":
            preset_id = "RGBD_camera"
    preset = get_builtin_preset(preset_id)
    if preset is None:
        raise HTTPException(status_code=500, detail="Wizard preset not found.")
    config = preset["config"]
    insights = compute_insights(config["preprocessPlugins"], config["reconstructionPlugin"])
    return {
        "presetId": preset_id,
        "title": preset["title"],
        "config": config,
        "stageCards": config_to_stage_cards(config),
        **insights,
    }


@app.get("/api/demo")
def demo(surfaceType: str = "Сфера", count: int = 350) -> dict[str, Any]:
    return demo_payload(surfaceType, count)


@app.get("/api/demo/types")
def demo_types() -> list[str]:
    return DEMO_SURFACE_TYPES


@app.get("/api/generator/catalog")
def generator_catalog() -> dict[str, Any]:
    return generator_catalog_payload()


@app.post("/api/generator/layer")
def generator_layer(body: GeneratorLayerBody) -> dict[str, Any]:
    try:
        return generate_layer(body.kind, body.type, body.params)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.post("/api/generator/resolve-intersections")
def generator_resolve_intersections(body: GeneratorResolveBody) -> dict[str, Any]:
    try:
        layers = [item.model_dump() for item in body.layers]
        resolved = resolve_intersections(
            layers,
            eps=body.eps,
            clip_surface_inside_objects=body.clipSurfaceInsideObjects,
            clip_objects_vs_objects=body.clipObjectsVsObjects,
        )
        return {
            "layers": resolved,
            "removedTotal": sum(int(layer.get("removedCount", 0)) for layer in resolved),
        }
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.post("/api/dataset/generate")
async def dataset_generate(
    count: int = Form(5),
    seed: int = Form(42),
    outputDir: str = Form("sonar_dataset"),
    objectScale: float = Form(1.0),
    beamCount: int = Form(45),
    lengthCount: int | None = Form(None),
    model: UploadFile = File(...),
) -> dict[str, Any]:
    filename = (model.filename or "").strip()
    if not filename.lower().endswith(".obj"):
        raise HTTPException(status_code=400, detail="Upload a .obj 3D model file.")
    try:
        raw = await model.read()
        text = raw.decode("utf-8", errors="ignore")
        object_points = load_object_points_from_obj_text(text)
        return generate_dataset(
            count=count,
            seed=seed,
            output_dir=outputDir or "sonar_dataset",
            object_points=object_points,
            object_name=filename,
            object_scale=objectScale,
            beam_count=beamCount,
            length_count=lengthCount,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except OSError as exc:
        raise HTTPException(status_code=500, detail=f"Failed to write dataset: {exc}") from exc


@app.post("/api/dataset/preview")
def dataset_preview(body: DatasetPreviewBody) -> dict[str, Any]:
    try:
        return load_scene_preview(
            stem=body.stem,
            output_dir=body.outputDir or "sonar_dataset",
            max_points=body.maxPoints,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except FileNotFoundError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    except OSError as exc:
        raise HTTPException(status_code=500, detail=f"Failed to load scene: {exc}") from exc


@app.post("/api/generator/export")
def generator_export(body: GeneratorExportBody) -> Response:
    from urllib.parse import quote

    fmt = (body.format or "xyz").lower().lstrip(".")
    if fmt not in ("xyz", "ply", "obj", "npy"):
        raise HTTPException(status_code=400, detail="Supported formats: xyz, ply, obj, npy")

    filename = body.filename
    content: bytes
    media: str
    ext: str

    if fmt == "npy":
        try:
            if body.layers:
                rows = layers_to_pointnet_rows([item.model_dump() for item in body.layers])
            elif body.points is not None:
                label = 0.0 if body.classLabel is None else float(body.classLabel)
                rows = points_to_pointnet_rows(body.points, label)
            else:
                raise HTTPException(status_code=400, detail="Provide points or layers to export.")
            content = export_npy_float64(rows)
        except HTTPException:
            raise
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        media = "application/octet-stream"
        ext = "npy"
    else:
        points = body.points
        if points is None and body.layers:
            try:
                points = merge_layers_world([item.model_dump() for item in body.layers])
            except ValueError as exc:
                raise HTTPException(status_code=400, detail=str(exc)) from exc
        if points is None:
            raise HTTPException(status_code=400, detail="Provide points or layers to export.")

        if fmt == "ply":
            text = export_ply(points)
            media = "application/octet-stream"
            ext = "ply"
        elif fmt == "obj":
            stem = Path(body.filename or "cloud").stem or "cloud"
            text = export_obj(points, object_name=stem)
            media = "text/plain; charset=utf-8"
            ext = "obj"
        else:
            text = export_xyz(points)
            media = "text/plain; charset=utf-8"
            ext = "xyz"
        content = text.encode("utf-8")

    out_name = filename or f"cloud.{ext}"
    if not out_name.lower().endswith(f".{ext}"):
        out_name = f"{out_name}.{ext}"

    ascii_name = "".join(ch if 32 <= ord(ch) < 127 and ch not in '\\/"' else "_" for ch in out_name)
    if not ascii_name.lower().endswith(f".{ext}"):
        ascii_name = f"cloud.{ext}"
    disposition = (
        f"attachment; filename=\"{ascii_name}\"; "
        f"filename*=UTF-8''{quote(out_name)}"
    )

    return Response(
        content=content,
        media_type=media,
        headers={"Content-Disposition": disposition},
    )


@app.get("/api/user-presets")
def get_user_presets() -> list[dict[str, Any]]:
    return load_user_presets()


@app.post("/api/user-presets")
def post_user_preset(body: UserPresetBody) -> dict[str, Any]:
    presets = load_user_presets()
    preset_id = body.idValue or f"user:{body.title.lower().replace(' ', '_')}"
    preset = {
        "title": body.title,
        "idValue": preset_id,
        "stages": body.stages,
    }
    replaced = False
    for index, existing in enumerate(presets):
        if existing.get("idValue") == preset_id or existing.get("title") == body.title:
            presets[index] = preset
            replaced = True
            break
    if not replaced:
        presets.append(preset)
    save_user_presets(presets)
    return preset


@app.post("/api/run")
async def run_pipeline(
    file: UploadFile | None = File(default=None),
    preset_id: str | None = Form(default=None),
    config_json: str | None = Form(default=None),
    demo_surface: str | None = Form(default=None),
    geometry_format: str = Form(default="json"),
) -> dict[str, Any]:
    if not DOTSTOSURFACE_BIN.is_file():
        raise HTTPException(status_code=500, detail=f"Binary not found: {DOTSTOSURFACE_BIN}")

    work_id = uuid.uuid4().hex
    work_dir = Path(tempfile.gettempdir()) / "dottosurface" / work_id
    work_dir.mkdir(parents=True, exist_ok=True)
    output_path = work_dir / "result.json"

    try:
        if demo_surface:
            demo = demo_payload(demo_surface)
            input_path = work_dir / "demo.xyz"
            lines = [f"{p[0]} {p[1]} {p[2]}" for p in demo["points"]]
            input_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
        elif file is not None:
            suffix = Path(file.filename or "cloud.ply").suffix.lower() or ".ply"
            raw = await file.read()
            if suffix == ".obj":
                try:
                    text = raw.decode("utf-8", errors="ignore")
                    obj_points = parse_obj_points(text)
                except Exception as exc:
                    raise HTTPException(status_code=400, detail=f"Invalid OBJ: {exc}") from exc
                if not obj_points:
                    raise HTTPException(status_code=400, detail="OBJ has no vertices (v x y z).")
                input_path = work_dir / "input.xyz"
                input_path.write_text(
                    "\n".join(f"{p[0]} {p[1]} {p[2]}" for p in obj_points) + "\n",
                    encoding="utf-8",
                )
            else:
                input_path = work_dir / f"input{suffix}"
                input_path.write_bytes(raw)
        else:
            raise HTTPException(status_code=400, detail="Provide file or demo_surface.")

        if config_json:
            pipeline_config = json.loads(config_json)
        elif preset_id:
            builtin = get_builtin_preset(preset_id)
            if builtin is not None:
                pipeline_config = builtin["config"]
            else:
                preset_path = next((p for p in list_preset_files() if p.stem == preset_id), None)
                if preset_path is None:
                    user_match = next(
                        (p for p in load_user_presets() if p.get("idValue") == preset_id),
                        None,
                    )
                    if user_match is None:
                        raise HTTPException(status_code=400, detail=f"Unknown preset: {preset_id}")
                    pipeline_config = preset_to_pipeline_config(user_match)
                else:
                    with preset_path.open("r", encoding="utf-8") as handle:
                        pipeline_config = preset_to_pipeline_config(json.load(handle))
        else:
            with DEFAULT_PIPELINE_CONFIG.open("r", encoding="utf-8") as handle:
                pipeline_config = json.load(handle)

        completed = run_pipeline_command(input_path, pipeline_config, output_path)
        if completed.returncode != 0:
            detail = completed.stderr.strip() or completed.stdout.strip() or "Pipeline failed."
            raise HTTPException(status_code=500, detail=detail)

        if not output_path.is_file():
            raise HTTPException(status_code=500, detail="Pipeline finished without output JSON.")

        with output_path.open("r", encoding="utf-8") as handle:
            result = json.load(handle)

        result = enrich_result(result, pipeline_config, completed.stdout.strip(), work_id)

        if geometry_format == "binary":
            bin_path = write_binary_geometry(work_dir, result)
            if bin_path:
                result["geometryUrl"] = f"/api/geometry/{work_id}"
                result.pop("points", None)
                result.pop("triangleIndices", None)

        return result
    except json.JSONDecodeError as exc:
        raise HTTPException(status_code=400, detail=f"Invalid config JSON: {exc}") from exc
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.get("/api/geometry/{work_id}")
def get_geometry(work_id: str) -> Response:
    bin_path = Path(tempfile.gettempdir()) / "dottosurface" / work_id / "geometry.bin"
    if not bin_path.is_file():
        raise HTTPException(status_code=404, detail="Geometry not found.")
    return Response(content=bin_path.read_bytes(), media_type="application/octet-stream")


@app.get("/airplane_reference.png")
def airplane_reference_image() -> FileResponse:
    candidates = [
        WEB_DIST / "airplane_reference.png",
        APP_ROOT / "assets" / "airplane_reference.png",
        APP_ROOT / "frontend" / "web" / "public" / "airplane_reference.png",
    ]
    for path in candidates:
        if path.is_file():
            return FileResponse(path, media_type="image/png")
    raise HTTPException(status_code=404, detail="Airplane reference image not found.")


@app.get("/")
def index() -> FileResponse:
    index_path = WEB_DIST / "index.html"
    if not index_path.is_file():
        raise HTTPException(
            status_code=503,
            detail="Frontend not built. Run: cd frontend/web && npm install && npm run build",
        )
    return FileResponse(index_path)


@app.get("/generator")
def generator_spa() -> FileResponse:
    return index()


@app.get("/dataset")
def dataset_spa() -> FileResponse:
    return index()


if (WEB_DIST / "assets").is_dir():
    app.mount("/assets", StaticFiles(directory=WEB_DIST / "assets"), name="assets")