from flask import Flask, request, jsonify
import cv2
import numpy as np
from pathlib import Path
import tempfile
import os
import logging
from hloc import logger as hloc_logger
from common.api import ImageMatchingAPI
from common.utils import (
    DEVICE,
    get_matcher_zoo,
    load_config,
    ROOT,
)

app = Flask(__name__)

# Load configuration and matcher zoo
config = load_config(ROOT / "common/config.yaml")
matcher_zoo_restored = get_matcher_zoo(config["matcher_zoo"])

# Initialize API instances for each model
api_instances = {}
for model_name, conf in matcher_zoo_restored.items():
    api_instances[model_name] = ImageMatchingAPI(conf=conf, device=DEVICE)

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def match_images(img_path1, img_path2, api_instance):
    logger.info(f"Matching images with {DEVICE} using provided API instance")
    try:
        # Read the images using OpenCV
        image0 = cv2.imread(img_path1)
        image1 = cv2.imread(img_path2)

        if image0 is None or image1 is None:
            raise ValueError("One or both images could not be read. Ensure the files are valid images.")

        # Convert BGR to RGB
        image0 = cv2.cvtColor(image0, cv2.COLOR_BGR2RGB)
        image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)

        results = api_instance(image0, image1)

        return {
            "num_keypoints0": len(results["keypoints0_orig"]),
            "num_keypoints1": len(results["keypoints1_orig"]),
            "num_matches": len(results["mkeypoints0_orig"]),
            "num_inliers": len(results["mmkeypoints0_orig"]),
            "keypoints0": results["keypoints0_orig"].tolist(),
            "keypoints1": results["keypoints1_orig"].tolist(),
            "matches0": results["mkeypoints0_orig"].tolist(),
            "matches1": results["mmkeypoints0_orig"].tolist(),
            "inliers0": results["mmkeypoints0_orig"].tolist(),
            "inliers1": results["mmkeypoints0_orig"].tolist(),
            "confidence": results["mconf"].tolist(),
            "inlier_confidence": results["mmconf"].tolist(),
        }
    except Exception as e:
        logger.error("Matching failed: %s", str(e))
        return {"error": str(e)}
    finally:
        try:
            os.remove(img_path1)
            os.remove(img_path2)
        except Exception as cleanup_error:
            logger.error("Failed to clean up temporary files: %s", str(cleanup_error))

@app.route('/match-images', methods=['POST'])
def match_images_endpoint():
    if 'img1' not in request.files or 'img2' not in request.files:
        return jsonify({"error": "Please upload both images."}), 400

    img1_file = request.files['img1']
    img2_file = request.files['img2']
    model_name = request.form.get('model', 'superpoint+superglue')

    # Validate file types
    valid_extensions = {'png', 'jpg', 'jpeg', 'bmp', 'tiff'}
    if not (img1_file.filename.split('.')[-1].lower() in valid_extensions and
            img2_file.filename.split('.')[-1].lower() in valid_extensions):
        return jsonify({"error": "Invalid file type. Please upload images in one of the following formats: png, jpg, jpeg, bmp, tiff"}), 400

    try:
        # Save the uploaded files to a temporary directory with unique filenames
        temp_dir = tempfile.mkdtemp()
        img_path1 = os.path.join(temp_dir, next(tempfile._get_candidate_names()) + os.path.splitext(img1_file.filename)[1])
        img_path2 = os.path.join(temp_dir, next(tempfile._get_candidate_names()) + os.path.splitext(img2_file.filename)[1])
        img1_file.save(img_path1)
        img2_file.save(img_path2)

        # Get the pre-initialized API instance for the selected model
        api_instance = api_instances.get(model_name)
        if not api_instance:
            raise ValueError(f"Model '{model_name}' not found in matcher zoo.")

        # Run the matching task synchronously
        result = match_images(img_path1, img_path2, api_instance)
    except Exception as e:
        logger.error("Error processing images: %s", str(e))
        return jsonify({"error": str(e)}), 500
    finally:
        # Clean up the temporary directory
        if os.path.exists(temp_dir):
            try:
                os.rmdir(temp_dir)
            except Exception as cleanup_error:
                logger.error("Failed to clean up temporary directory: %s", str(cleanup_error))

    return jsonify(result), 200

if __name__ == '__main__':
    app.run(debug=True)