import os
import sys
import gc
import base64
import tempfile
import traceback
import io
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS

# Import the logic from the backend modules
from centerline_extraction import compute_centerlines_network, load_nrrd_as_vtk_image
from ostium_detection import extract_centerline_data
from endograft_generation import build_endograft_stl_payload
from maquette_generation import generate_fenestrated_maquette
from endograft_selection import start_endograft_pipeline

# Flask API server that wraps the VMTK centerline extraction, ostium detection, and endograft generation pipelines
app = Flask(__name__)
CORS(app)  # Allow cross-origin requests from the Next.js frontend
app.config["MAX_CONTENT_LENGTH"] = 200 * 1024 * 1024 # Maximum file upload size (200 MB)

# Simple health check so the deployment platform knows the server is running
@app.route("/health", methods=["GET"])
def health():
  return jsonify({"status": "ok"})

# Runs the full VMTK centerline extraction pipeline on NRRD segmentation and returns the result as JSON
@app.route("/extract", methods=["POST"])
def extract():
  # Check that the segmentation file sent to the centerline extraction endpoint exists and is an NRRD file
  if "file" not in request.files:
    return jsonify({"status": "error", "error": "No file provided"}), 400

  uploaded = request.files["file"]
  file_name = uploaded.filename.lower()
  if not file_name.endswith(".nrrd"):
    return jsonify({"status": "error", "error": "Invalid File Type. Allowed: .nrrd"}), 400

  # Save to a temp file so VMTK can read it from disk
  tmp = tempfile.NamedTemporaryFile(suffix=".nrrd", delete=False)
  try:
    uploaded.save(tmp.name)
    tmp.close()

    # Read the NRRD segmentation via SimpleITK (preserves full direction matrix)
    image, origin, direction = load_nrrd_as_vtk_image(tmp.name)

    # Voronoi centerline extraction
    centerlines, branch_junctions = compute_centerlines_network(image, origin, direction)

    # Extract labelled segments and artery bifurcation data
    data = extract_centerline_data(centerlines, branch_junctions=branch_junctions)

    # Free the large VTK pipeline objects immediately
    del centerlines, branch_junctions
    gc.collect()

    # Include the NRRD image's physical extent so the frontend can normalize the centerline
    dims = image.GetDimensions()
    spacing = image.GetSpacing()
    phys_size = [dims[i] * spacing[i] for i in range(3)]
    phys_center = [origin[i] + phys_size[i] / 2.0 for i in range(3)]
    data['volume_info'] = {'center': {'x': phys_center[0], 'y': phys_center[1], 'z': phys_center[2]}, 'size': {'x': phys_size[0], 'y': phys_size[1], 'z': phys_size[2]}}

    del image
    gc.collect()

    return jsonify({"status": "success", "data": data})

  except TimeoutError as e:
    return jsonify({"status": "error", "error": str(e)}), 504

  except Exception as e:
    traceback.print_exc()
    return jsonify({"status": "error", "error": str(e)}), 500

  finally:
    # Always clean up the temp file
    try:
      os.unlink(tmp.name)
    except OSError:
      pass

# Runs the endograft selection pipeline on computed bounds of aneurysm
@app.route("/select_graft", methods=["POST"])
def select_graft():
  body = request.get_json(silent=True) or {}
  proximal_diameter = body.get("proximal_diameter", 0)
  distal_diameter = body.get("distal_diameter", 0)
  proximal_start_point = body.get("proximal_start_point")
  distal_end_point = body.get("distal_end_point")
  length = body.get("length", 0)
  fenestrations_exist = body.get("fenestrations_exist", False)
  distance_to_last_fenestration = body.get("distance_to_last_fenestration", 0)

  try:
    selected_graft = start_endograft_pipeline(
      proximal_diameter, distal_diameter, proximal_start_point, distal_end_point, length, fenestrations_exist, distance_to_last_fenestration
    )
    return jsonify({"status": "success", "data": selected_graft})
  except Exception as e:
    traceback.print_exc()
    return jsonify({"status": "error", "error": str(e)}), 500

# Runs the endograft generation pipeline on VMTK centerline and returns the result as JSON
@app.route("/endograft", methods=["POST"])
def endograft():
  body = request.get_json(silent=True) or {}
  centerline = body.get("centerline")
  params = body.get("params") or {}

  if not isinstance(centerline, list) or len(centerline) < 2:
    return jsonify({"status": "error", "error": "centerline must be a list of at least 2 points"}), 400

  try:
    payload = build_endograft_stl_payload(centerline, include_combined=False, **params)
    result = {"endograft_stl_b64": base64.b64encode(payload["endograft"]).decode("utf-8"), "struts_stl_b64": base64.b64encode(payload["struts"]).decode("utf-8")}
    return jsonify({"status": "success", "data": result})
  except Exception as e:
    traceback.print_exc()
    return jsonify({"status": "error", "error": str(e)}), 500

# Generates endograft STL maquette
@app.route("/maquette", methods=["POST"])
def maquette():
  body = request.get_json(silent=True) or {}
  centerline = body.get("centerline")
  params = body.get("params") or {}
  fenestrations = body.get("fenestrations") or []
  include_struts = body.get("include_struts", True)

  if not isinstance(centerline, list) or len(centerline) < 2:
    return jsonify({"status": "error", "error": "centerline must be a list of at least 2 points"}), 400

  try:
    stl_bytes = generate_fenestrated_maquette(centerline, params, fenestrations, include_struts)
    result = {"maquette_stl_b64": base64.b64encode(stl_bytes).decode("utf-8")}
    return jsonify({"status": "success", "data": result})
  except Exception as e:
    traceback.print_exc()
    return jsonify({"status": "error", "error": str(e)}), 500


if __name__ == "__main__":
  port = int(os.environ.get("PORT", 5050))
  app.run(host="0.0.0.0", port=port)