from flask import Flask, render_template, request, redirect, url_for, send_from_directory, session
import json
import random
import os
import string
import logging
from datetime import datetime

# Set up logging
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    handlers=[
                        logging.FileHandler("app.log"),
                        logging.StreamHandler()
                    ])
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.config['SECRET_KEY'] = 'supersecretkey'  # Change this to a random secret key

# Directories for visualizations
VISUALIZATION_DIRS_PLAN_OF_SQLS = {
    "TP": "htmls_POS/TP",
    "TN": "htmls_POS/TN",
    "FP": "htmls_POS/FP",
    "FN": "htmls_POS/FN"
}

VISUALIZATION_DIRS_CHAIN_OF_TABLE = {
    "TP": "htmls_COT/TP",
    "TN": "htmls_COT/TN",
    "FP": "htmls_COT/FP",
    "FN": "htmls_COT/FN"
}


def save_session_data(username, data):
    try:
        base_dir = os.path.dirname(os.path.abspath(__file__))
        session_dir = os.path.join(base_dir, 'session_data')
        os.makedirs(session_dir, exist_ok=True)

        file_path = os.path.join(session_dir, f'{username}_session.json')

        with open(file_path, 'w') as f:
            json.dump(data, f, indent=4)

        logger.info(f"Session data saved for user {username} at {file_path}")
    except Exception as e:
        logger.exception(f"Error saving session data for user {username}: {e}")


# Similarly, update the load_session_data function
def load_session_data(username):
    try:
        base_dir = os.path.dirname(os.path.abspath(__file__))
        file_path = os.path.join(base_dir, 'session_data', f'{username}_session.json')

        with open(file_path, 'r') as f:
            data = json.load(f)

        logger.info(f"Session data loaded for user {username} from {file_path}")
        return data
    except FileNotFoundError:
        logger.warning(f"No session data found for user {username}")
        return None
    except Exception as e:
        logger.exception(f"Error loading session data for user {username}: {e}")
        return None

# Load all sample files from the directories based on the selected method
def load_samples(method):
    logger.info(f"Loading samples for method: {method}")
    if method == "Chain-of-Table":
        visualization_dirs = VISUALIZATION_DIRS_CHAIN_OF_TABLE
    else:
        visualization_dirs = VISUALIZATION_DIRS_PLAN_OF_SQLS

    samples = {"TP": [], "TN": [], "FP": [], "FN": []}
    for category, dir_path in visualization_dirs.items():
        try:
            for filename in os.listdir(dir_path):
                if filename.endswith(".html"):
                    samples[category].append(filename)
            logger.info(f"Loaded {len(samples[category])} samples for category {category}")
        except Exception as e:
            logger.exception(f"Error loading samples from {dir_path}: {e}")
    return samples

# Randomly select balanced samples
def select_balanced_samples(samples):
    try:
        tp_fp_samples = random.sample(samples["TP"] + samples["FP"], 5)
        tn_fn_samples = random.sample(samples["TN"] + samples["FN"], 5)
        logger.info(f"Selected balanced samples: {len(tp_fp_samples + tn_fn_samples)}")
        return tp_fp_samples + tn_fn_samples
    except Exception as e:
        logger.exception("Error selecting balanced samples")
        return []

def generate_random_string(length=8):
    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))


@app.route('/', methods=['GET', 'POST'])
def index():
    logger.info("Rendering index page.")
    if request.method == 'POST':
        username = request.form.get('username')
        seed = request.form.get('seed')

        if not username or not seed:
            logger.error("Missing username or seed.")
            return "Missing username or seed", 400

        try:
            seed = int(seed)
            random.seed(seed)

            # Use only one method (e.g., "Chain-of-Table")
            method = "Chain-of-Table"
            all_samples = load_samples(method)
            selected_samples = select_balanced_samples(all_samples)
            logger.info(f"Number of selected samples: {len(selected_samples)}")

            if len(selected_samples) == 0:
                logger.error("No samples were selected.")
                return "No samples were selected", 500

            filename = f'{username}_{seed}_{method}_{generate_random_string()}.json'
            logger.info(f"Generated filename: {filename}")

            # Save session data
            session_data = {
                'responses': [],
                'username': username,
                'selected_samples': selected_samples,
                'method': method,
                'filename': filename,
                'current_index': 0
            }
            save_session_data(username, session_data)
            logger.info(f"Session data saved for user: {username}")

            return redirect(url_for('experiment', username=username))
        except Exception as e:
            logger.exception(f"Error in index route: {e}")
            return "An error occurred", 500
    return render_template('index.html')
@app.route('/experiment/<username>', methods=['GET', 'POST'])
def experiment(username):
    try:
        session_data = load_session_data(username)
        if not session_data:
            logger.error(f"No session data found for user: {username}")
            return redirect(url_for('index'))

        selected_samples = session_data['selected_samples']
        method = session_data['method']
        current_index = session_data['current_index']

        if current_index >= len(selected_samples):
            return redirect(url_for('completed', username=username))

        visualization_file = selected_samples[current_index]

        vis_dir = 'htmls_COT' if method == "Chain-of-Table" else 'htmls_POS'

        # Determine the correct visualization directory based on the category
        for category, dir_path in VISUALIZATION_DIRS_CHAIN_OF_TABLE.items():
            if visualization_file in os.listdir(dir_path):
                visualization_path = os.path.join(vis_dir, category, visualization_file)
                break
        else:
            logger.error(f"Visualization file {visualization_file} not found.")
            return "Visualization file not found", 404

        logger.info(f"Rendering experiment page with visualization: {visualization_path}")

        statement = """
Based on the explanation provided, what do you think the AI model will predict?
Will it predict the statement as TRUE or FALSE?
        """

        return render_template('experiment.html',
                               sample_id=current_index,
                               statement=statement,
                               visualization=url_for('send_visualization', filename=visualization_path),
                               username=username)
    except Exception as e:
        logger.exception(f"An error occurred in the experiment route: {e}")
        return "An error occurred", 500

@app.route('/feedback', methods=['POST'])
def feedback():
    try:
        username = request.form['username']
        prediction = request.form['prediction']

        session_data = load_session_data(username)
        if not session_data:
            logger.error(f"No session data found for user: {username}")
            return redirect(url_for('index'))

        # Store the user's prediction
        session_data['responses'].append({
            'sample_id': session_data['current_index'],
            'user_prediction': prediction
        })

        # Move to the next sample
        session_data['current_index'] += 1

        # Save updated session data
        save_session_data(username, session_data)
        logger.info(f"Prediction saved for user {username}, sample {session_data['current_index'] - 1}")

        if session_data['current_index'] >= len(session_data['selected_samples']):
            return redirect(url_for('completed', username=username))

        return redirect(url_for('experiment', username=username))
    except Exception as e:
        logger.exception(f"Error in feedback route: {e}")
        return "An error occurred", 500


@app.route('/completed/<username>')
def completed(username):
    try:
        session_data = load_session_data(username)
        if not session_data:
            logger.error(f"No session data found for user: {username}")
            return redirect(url_for('index'))

        responses = session_data['responses']
        method = session_data['method']

        json_file = 'Tabular_LLMs_human_study_vis_6_COT.json' if method == "Chain-of-Table" else 'Tabular_LLMs_human_study_vis_6_POS.json'

        with open(json_file, 'r') as f:
            ground_truth = json.load(f)

        correct_predictions = 0
        true_predictions = 0
        false_predictions = 0

        for response in responses:
            sample_id = response['sample_id']
            user_prediction = response['user_prediction']
            visualization_file = session_data['selected_samples'][sample_id]
            index = visualization_file.split('-')[1].split('.')[0]  # Extract index from filename

            ground_truth_key = f"COT_test-{index}.html" if method == "Chain-of-Table" else f"POS_test-{index}.html"

            if ground_truth_key in ground_truth:
                model_prediction = ground_truth[ground_truth_key]['answer'].upper()
                if user_prediction.upper() == model_prediction:
                    correct_predictions += 1

                if user_prediction.upper() == "TRUE":
                    true_predictions += 1
                elif user_prediction.upper() == "FALSE":
                    false_predictions += 1
            else:
                logger.warning(f"Missing key in ground truth: {ground_truth_key}")

        accuracy = (correct_predictions / len(responses)) * 100 if responses else 0
        accuracy = round(accuracy, 2)

        true_percentage = (true_predictions / len(responses)) * 100 if len(responses) else 0
        false_percentage = (false_predictions / len(responses)) * 100 if len(responses) else 0

        true_percentage = round(true_percentage, 2)
        false_percentage = round(false_percentage, 2)

        return render_template('completed.html',
                               accuracy=accuracy,
                               true_percentage=true_percentage,
                               false_percentage=false_percentage)
    except Exception as e:
        logger.exception(f"An error occurred in the completed route: {e}")
        return "An error occurred", 500


@app.route('/visualizations/<path:filename>')
def send_visualization(filename):
    logger.info(f"Attempting to serve file: {filename}")
    # Ensure the path is safe and doesn't allow access to files outside the intended directory
    base_dir = os.getcwd()
    file_path = os.path.normpath(os.path.join(base_dir, filename))
    if not file_path.startswith(base_dir):
        return "Access denied", 403

    if not os.path.exists(file_path):
        return "File not found", 404

    directory = os.path.dirname(file_path)
    file_name = os.path.basename(file_path)
    logger.info(f"Serving file from directory: {directory}, filename: {file_name}")
    return send_from_directory(directory, file_name)


if __name__ == "__main__":
    os.makedirs('session_data', exist_ok=True)  # Ensure the directory for session files exists
    app.run(host="0.0.0.0", port=7860, debug=True)