from flask import Flask, render_template, request, jsonify, send_file
from werkzeug.utils import secure_filename
from datetime import datetime
import os, subprocess, shutil, json

app = Flask(__name__)
ALLOWED_EXTENSIONS = {'jpg'}


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

# Run the prediction script
def process_image(file_path):
    try:
        result = subprocess.check_output(["python", "tools/predict.py", "--cfg", "experiments/cls_hrnet_w48_sgd_lr5e-2_wd1e-4_bs32_x100.yaml", "--testModel", "output/imagenet/cls_hrnet_w48_sgd_lr5e-2_wd1e-4_bs32_x100/model_best.pth.tar", "--image_path", file_path], stderr=subprocess.STDOUT)
        return result
    except subprocess.CalledProcessError as e:
        print(f"Error: {e.output}")
        return None


# Static Files
@app.route('/images/<filename>')
def images(filename):
    return send_file(os.path.join(app.root_path, 'templates', 'images', filename), mimetype='image/jpeg')



# Home Page
@app.route('/')
def index():
    with open(os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.json"), 'r') as json_file:
        log_data = json.load(json_file)
    log_data = sorted(log_data, key=lambda x: x['date_time'], reverse=True)
    return render_template('home.html', log_data=log_data)


# Upload Process - Save Image - Run Prediction - Save Log
@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400
    file = request.files['file']
    filename = secure_filename(file.filename)

    if allowed_file(filename):
        file_path = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "uploaded.jpg")
        file.save(file_path)
        result = process_image(file_path)
        if result is not None:
            output_file_path_txt = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.txt")
            output_file_path_json = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.json")
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            with open(output_file_path_txt, 'a') as output_file:
                output_file.write(f"{current_time} {filename}: {result.decode('utf-8')}")
            with open(output_file_path_json, 'r') as json_file:
                log_data = json.load(json_file)
            log_data.append({
                'date_time': current_time,
                'filename': filename,
                'class_accuracy': result.decode('utf-8').strip(),
                'expected_class': 'Not Provided'
            })
            with open(output_file_path_json, 'w') as json_file:
                json.dump(log_data, json_file)
            # Use current time as a unique query parameter to append to the image URL
            unique_query = datetime.now().strftime("%Y%m%d%H%M%S")
            img_url = f"/static/images/uploaded_images/uploaded.jpg?{unique_query}"
            return jsonify({'result': result.decode('utf-8'), 'file_type': 'jpg', 'img_url': img_url})
        else:
            return jsonify({'error': 'Error processing image. Please try again.'}), 500
    else:
        return jsonify({'error': 'Invalid file format. Please upload a jpg file.'}), 400






# All Log Processes
# Log Process - Append Choice
@app.route('/append_choice', methods=['POST'])
def append_choice():
    choice = request.json.get('choice')
    if not choice:
        return jsonify({'error': 'No choice provided'}), 400
    output_file_path_txt = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.txt")
    output_file_path_json = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.json")
    with open(output_file_path_txt, 'a') as output_file:
        output_file.write(f"Verified Class: {choice}\n")
    with open(output_file_path_json, 'r') as json_file:
        log_data = json.load(json_file)
    log_data[-1]['expected_class'] = choice
    with open(output_file_path_json, 'w') as json_file:
        json.dump(log_data, json_file)
    return jsonify({'message': 'Choice appended successfully'})


# Log Process - Download Log
@app.route('/download')
def download_file():
    output_file_path_txt = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.txt")
    return send_file(output_file_path_txt, as_attachment=True, download_name="log.txt")

# Log Process - Get Log Data
@app.route('/get_log_data')
def get_log_data():
    output_file_path_json = os.path.join(app.root_path, 'static', 'images', 'uploaded_images', "log.json")
    with open(output_file_path_json, 'r') as json_file:
        log_data = json.load(json_file)
    return jsonify(log_data)

# Log Process - Clear Log
@app.route('/clear_log', methods=['POST'])
def clear_log():
    log_json_path = os.path.join('static', 'images', 'uploaded_images', 'log.json')
    log_txt_path = os.path.join('static', 'images', 'uploaded_images', 'log.txt')
    
    with open(log_json_path, 'w') as log_file:
        log_file.write('[]') 
    if os.path.exists(log_txt_path):
        with open(log_txt_path, 'w') as log_file:
            log_file.write('')  
    return 'Log cleared', 200






# About Page
@app.route('/about')
def about():
    return render_template('about.html')


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