from flask import Flask, request, jsonify, render_template
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.data import MetadataCatalog
from detectron2.utils.visualizer import Visualizer, ColorMode
import numpy as np
from PIL import Image
import io
import os
import requests
import gdown

# Initialize Flask app
app = Flask(__name__)
cfg = None
# Google Drive file URL
GDRIVE_MODEL_URL = "https://drive.google.com/uc?id=18aEDo-kWOBhg8mAhnbpFkuM6bmmrBH4E"  # Replace 'your-file-id' with the actual file ID from Google Drive
LOCAL_MODEL_PATH = "model_final.pth"
def download_file_from_google_drive(id, destination):
    gdown.download(GDRIVE_MODEL_URL, LOCAL_MODEL_PATH, quiet=False)

                
file_id = '18aEDo-kWOBhg8mAhnbpFkuM6bmmrBH4E'
destination = 'model_final.pth'
download_file_from_google_drive(file_id, destination)

# Download model from Google Drive if not already present locally
def download_model():
    if not os.path.exists(LOCAL_MODEL_PATH):
        response = requests.get(GDRIVE_MODEL_URL, stream=True)
        if response.status_code == 200:
            with open(LOCAL_MODEL_PATH, 'wb') as f:
                f.write(response.content)
        else:
            raise Exception(f"Failed to download model from Google Drive: {response.status_code}")

# Configuration and model setup
def setup_model(model_path):
    global cfg
    cfg = get_cfg()
    cfg.merge_from_file("config.yaml")  # Update with the config file path
    cfg.MODEL.WEIGHTS = model_path
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
    cfg.MODEL.DEVICE = "cpu"  # Use "cuda" for GPU
    return DefaultPredictor(cfg)

# Ensure model is available
predictor = setup_model(LOCAL_MODEL_PATH)

# Define expected parts and costs
expected_parts = ['headlamp', 'rear_bumper', 'door', 'hood', 'front_bumper']
cost_dict = {
    'headlamp': 300,
    'rear_bumper': 250,
    'door': 200,
    'hood': 220,
    'front_bumper': 250,
    'other': 150
}

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload():
    if 'file' not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No file selected"}), 400

    # Load image
    image = Image.open(file).convert("RGB")
    image_np = np.array(image)

    # Run model prediction
    outputs = predictor(image_np)
    instances = outputs["instances"].to("cpu")
    class_names = MetadataCatalog.get(cfg.DATASETS.TEST[0]).thing_classes

    # Initialize total cost
    total_cost = 0
    damage_details = []

    for j in range(len(instances)):
        class_id = instances.pred_classes[j].item()
        damaged_part = class_names[class_id] if class_id < len(class_names) else 'unknown'
        if damaged_part not in expected_parts:
            damaged_part = 'other'

        repair_cost = cost_dict.get(damaged_part, cost_dict['other'])
        total_cost += repair_cost

        damage_details.append({
            'part': damaged_part,
            'cost_usd': repair_cost
        })

    response = {
        "damages": damage_details,
        "total_cost": total_cost
    }

    return jsonify(response)

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