from flask import Blueprint, request, jsonify
from werkzeug.utils import secure_filename
from models.user import User
from models.image import Image
from mongoengine.errors import DoesNotExist
import os
import subprocess
import json
import time
from utils.jwt_service import jwt_required

upload_bp = Blueprint('upload_bp', __name__)

# Allowed settings
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
ALLOWED_MIME_TYPES = {'image/png', 'image/jpeg'}

# Folder paths
UPLOAD_FOLDER = os.path.join('public', 'images', 'uploads')
OUTPUT_FOLDER = os.path.join('public', 'images', 'output_files')

# Ensure folders exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

def allowed_file(file):
    return (
        '.' in file.filename and
        file.mimetype in ALLOWED_MIME_TYPES and
        file.filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
    )

@upload_bp.route('/upload', methods=['POST'])
@jwt_required
def upload_images():
    userid = getattr(request, "user_id", None)
    # userid = request.form.get('userid')
    if not userid:
        return jsonify({'error': 'User ID is required'}), 400

    try:
        user = User.objects.get(id=userid)
    except DoesNotExist:
        return jsonify({'error': 'User not found'}), 404

    files = request.files.getlist('images')
    if not files or not isinstance(files, list):
        return jsonify({'status': False, 'message': 'No images provided in the request.'}), 400

    image_urls, image_names = [], []

    for file in files:
        if file and allowed_file(file):
            filename = secure_filename(file.filename)
            unique_name = f"{int(time.time())}_{filename}"
            save_path = os.path.join(UPLOAD_FOLDER, unique_name)
            file.save(save_path)
            image_urls.append(os.path.abspath(save_path))
            image_names.append(unique_name)
        else:
            return jsonify({'status': False, 'message': 'Invalid file type. Only JPG, PNG are allowed'}), 400

    if not image_urls:
        return jsonify({'status': False, 'message': 'No valid images uploaded.'}), 400

    # Save initial record
    image_record = Image(
        user_id=userid,
        image=",".join(image_names)
    ).save()

    output_filename = f"extracted_text_{userid}_{str(image_record.id)}.json"
    output_file_path = os.path.join(OUTPUT_FOLDER, output_filename)

    success, extracted_text, error = run_python_script(image_urls, output_file_path)
    if not success:
        return jsonify({'status': False, 'message': 'Error executing OCR script.', 'details': error}), 500

    # Update the image record
    image_record.update(
        output_file=output_file_path,
        extracted_text=extracted_text
    )

    return jsonify({
        'status': True,
        'message': 'Image uploaded and text extracted successfully.',
        'output_file': output_file_path,
        'extracted_text': extracted_text
    })


def run_python_script(image_paths, output_file):
    python_path = '/home/ubuntu/myvenv/bin/python'  # Update if different
    script_path = os.path.join('ocr', 'OCR.py')
    model_path = os.path.join('ocr', 'Obj_detector_model.pt')

    command = [
        python_path, script_path,
        '--model_path', model_path,
        '--images', *image_paths,
        '--output_file', output_file
    ]

    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        with open(output_file, 'r') as f:
            extracted_text = json.load(f)
        return True, extracted_text, ""
    except subprocess.CalledProcessError as e:
        return False, None, e.stderr or e.stdout
    except Exception as e:
        return False, None, str(e)