import os
import cv2
import numpy as np
from flask import Flask, request, jsonify, send_file
from werkzeug.utils import secure_filename

app = Flask(__name__)

# ==============================================================================
# [CONFIG: STORAGE DIRECTORIES]
# ------------------------------------------------------------------------------
# Modify these paths to change where incoming raw uploads and processed output
# images are stored on your server's disk.
# ==============================================================================
UPLOAD_FOLDER = os.path.join(os.getcwd(), "storage", "inputs")   # <--- INPUT DIRECTORY
OUTPUT_FOLDER = os.path.join(os.getcwd(), "storage", "outputs")  # <--- OUTPUT DIRECTORY

app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config["OUTPUT_FOLDER"] = OUTPUT_FOLDER

# Ensure local directories exist on startup
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)


def align_image_to_reference(ref_path: str, target_path: str, output_path: str):
    """Aligns target_path image to match ref_path image geometry and saves to output_path."""
    ref_img = cv2.imread(ref_path)
    targ_img = cv2.imread(target_path)

    if ref_img is None or targ_img is None:
        raise ValueError("Could not read input images from storage.")

    # Convert to grayscale for SIFT feature extraction
    gray_ref = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
    gray_targ = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY)

    # 1. Detect SIFT features
    sift = cv2.SIFT_create()
    kp_ref, des_ref = sift.detectAndCompute(gray_ref, None)
    kp_targ, des_targ = sift.detectAndCompute(gray_targ, None)

    if des_ref is None or des_targ is None:
        raise ValueError("Failed to extract keypoints from one or both images.")

    # 2. Match keypoints using FLANN
    INDEX_KDTREE = 1
    flann = cv2.FlannBasedMatcher(
        dict(algorithm=INDEX_KDTREE, trees=5), dict(checks=50)
    )
    matches = flann.knnMatch(des_ref, des_targ, k=2)

    # 3. Apply Lowe's ratio test to filter matches
    good_matches = [m for m, n in matches if m.distance < 0.7 * n.distance]

    if len(good_matches) < 10:
        raise ValueError("Insufficient matching features found between images.")

    # 4. Extract keypoint coordinates
    src_pts = np.float32([kp_ref[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    dst_pts = np.float32([kp_targ[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

    # 5. Compute Affine Transformation Matrix (rigid: rotation, scale, translation)
    matrix, _ = cv2.estimateAffinePartial2D(dst_pts, src_pts, method=cv2.RANSAC)

    if matrix is None:
        raise ValueError("Failed to compute valid alignment transformation matrix.")

    # 6. Warp target image to match reference frame dimensions
    h, w = ref_img.shape[:2]
    aligned_img = cv2.warpAffine(
        targ_img,
        matrix,
        (w, h),
        flags=cv2.INTER_LANCZOS4,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(0, 0, 0)
    )

    # ==============================================================================
    # [OUTPUT IMAGE STORAGE LOCATION - WRITE TO DISK]
    # ------------------------------------------------------------------------------
    # The rotated/aligned image is saved here to output_path
    # ==============================================================================
    cv2.imwrite(output_path, aligned_img)


@app.route("/align", methods=["POST"])
def align_endpoint():
    # Validate request payload
    if "reference" not in request.files or "target" not in request.files:
        return jsonify({"error": "Missing 'reference' or 'target' file in request form-data."}), 400

    ref_file = request.files["reference"]
    target_file = request.files["target"]

    if ref_file.filename == "" or target_file.filename == "":
        return jsonify({"error": "No file selected."}), 400

    # Sanitize filenames
    ref_name = secure_filename(ref_file.filename)
    target_name = secure_filename(target_file.filename)

    # ==============================================================================
    # [INPUT IMAGE STORAGE LOCATION - SAVE RECEIVED FILES]
    # ------------------------------------------------------------------------------
    # Input files are saved into 'app.config["UPLOAD_FOLDER"]'
    # ==============================================================================
    input_ref_path = os.path.join(app.config["UPLOAD_FOLDER"], f"ref_{ref_name}")
    input_target_path = os.path.join(app.config["UPLOAD_FOLDER"], f"target_{target_name}")

    ref_file.save(input_ref_path)       # <-- Input Reference saved here
    target_file.save(input_target_path) # <-- Input Target saved here
    # ==============================================================================

    # ==============================================================================
    # [OUTPUT IMAGE STORAGE LOCATION - DEFINE TARGET PATH]
    # ------------------------------------------------------------------------------
    # Rotated image path in 'app.config["OUTPUT_FOLDER"]'
    # ==============================================================================
    output_aligned_path = os.path.join(app.config["OUTPUT_FOLDER"], f"aligned_{target_name}")
    # ==============================================================================

    try:
        # Run alignment pipeline
        align_image_to_reference(input_ref_path, input_target_path, output_aligned_path)

        # Return the processed image directly in the response
        return send_file(output_aligned_path, mimetype="image/png")

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


if __name__ == "__main__":
    # Run API server on http://0.0.0.0:5000
    app.run(host="0.0.0.0", port=5000, debug=True)