import os
import traceback
import uuid

import torch
from dotenv import load_dotenv
from flask import Flask, jsonify, render_template, request
from flask_cors import CORS
from ultralytics import YOLO
from werkzeug.utils import secure_filename

from processing import process_images


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
load_dotenv(os.path.join(BASE_DIR, '.env'))

app = Flask(__name__)
CORS(app)

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


def resolve_app_path(configured_path, default_relative_path):
    """Resolve configured relative paths from the application directory."""
    path = configured_path or default_relative_path
    if os.path.isabs(path):
        return path
    return os.path.join(BASE_DIR, path)


UPLOAD_FOLDER = resolve_app_path(os.getenv('UPLOAD_FOLDER'), os.path.join('static', 'uploads'))
MODELS_FOLDER = resolve_app_path(os.getenv('MODELS_FOLDER'), 'models')
PARTS_MODEL_NAME = os.getenv('PARTS_MODEL_NAME', 'best_parts_EP336.pt')
DAMAGE_MODEL_NAME = os.getenv('DAMAGE_MODEL_NAME', 'best_new_EP382.pt')


def resolve_model_path(model_name):
    """Resolve an absolute model path or a filename relative to MODELS_FOLDER."""
    if os.path.isabs(model_name):
        return model_name
    return os.path.join(MODELS_FOLDER, model_name)


PARTS_MODEL_PATH = resolve_model_path(PARTS_MODEL_NAME)
DAMAGE_MODEL_PATH = resolve_model_path(DAMAGE_MODEL_NAME)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(MODELS_FOLDER, exist_ok=True)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')


def load_model(model_path, label):
    """Load a YOLO model if its configured file exists."""
    if not os.path.isfile(model_path):
        print(f'Warning: {label} model file not found at {model_path}')
        return None

    try:
        model = YOLO(model_path)
        model.to(device)
        print(f"Successfully loaded {label.lower()} model '{os.path.basename(model_path)}' on {device}.")
        return model
    except Exception as exc:
        print(f'Error loading {label} model ({model_path}): {exc}')
        traceback.print_exc()
        return None


parts_model = load_model(PARTS_MODEL_PATH, 'Parts')
damage_model = load_model(DAMAGE_MODEL_PATH, 'Damage')


def allowed_file(filename):
    """Return True when filename has an allowed image extension."""
    return bool(filename) and '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def model_readiness_error():
    """Return a useful model readiness payload, or None when both models are loaded."""
    missing = []
    if parts_model is None:
        missing.append({'model': 'parts', 'path': PARTS_MODEL_PATH})
    if damage_model is None:
        missing.append({'model': 'damage', 'path': DAMAGE_MODEL_PATH})

    if not missing:
        return None

    return {
        'error': 'Prediction models are not ready.',
        'missing_or_unloaded_models': missing,
        'hint': 'Place the configured .pt files at the listed paths or update MODELS_FOLDER/model names in .env.'
    }


@app.route('/')
def home():
    """Serve the main HTML page."""
    return render_template('index.html')


@app.route('/health', methods=['GET'])
def health():
    """Expose basic service/model readiness without running inference."""
    readiness = model_readiness_error()
    if readiness:
        return jsonify({'status': 'degraded', **readiness}), 503
    return jsonify({'status': 'ok', 'device': device}), 200


@app.route('/predict', methods=['POST'])
def predict():
    """Receive one or more images and return part/damage predictions."""
    readiness = model_readiness_error()
    if readiness:
        return jsonify(readiness), 503

    raw_session_key = request.form.get('session_key', '').strip()
    session_key = secure_filename(raw_session_key)[:100] or uuid.uuid4().hex

    if 'file' not in request.files:
        return jsonify({'error': 'No file part in the request'}), 400

    files = request.files.getlist('file')
    if not files or all(not file.filename for file in files):
        return jsonify({'error': 'No selected files'}), 400

    saved_images = []
    skipped_files = []

    try:
        for index, file in enumerate(files):
            original_filename = file.filename or ''
            if not (file and allowed_file(original_filename)):
                skipped_files.append(original_filename or f'file_{index}')
                continue

            safe_original = secure_filename(original_filename) or f'image_{index}.jpg'
            unique_filename = f'{uuid.uuid4().hex}_{safe_original}'
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(filepath)
            saved_images.append({
                'path': filepath,
                'filename': original_filename,
                'index': index,
            })

        if not saved_images:
            return jsonify({
                'error': 'No valid files were uploaded. Allowed types: png, jpg, jpeg',
                'skipped_files': skipped_files,
            }), 400

        print(f"Processing {len(saved_images)} file(s) for session '{session_key}'...")
        results = process_images(parts_model, damage_model, saved_images)
        print(f"Processing complete for session '{session_key}'.")
        return jsonify(results)

    except Exception as exc:
        print(f'An error occurred during processing for session {session_key}: {exc}')
        traceback.print_exc()
        return jsonify({'error': f'An error occurred during processing: {exc}'}), 500

    finally:
        for image in saved_images:
            filepath = image['path']
            try:
                if os.path.exists(filepath):
                    os.remove(filepath)
            except Exception as exc:
                print(f'Error cleaning up file {filepath}: {exc}')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)