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
from huggingface_hub import login, HfApi, hf_hub_download

# 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__)

# Use the Hugging Face token from environment variables
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
else:
    logger.error("HF_TOKEN not found in environment variables")

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

# Directories for visualizations
VISUALIZATION_DIRS = {
    "No-XAI": "htmls_NO_XAI",
    "Dater": "htmls_DATER",
    "Chain-of-Table": "htmls_COT",
    "Plan-of-SQLs": "htmls_POS"
}

METHODS = ["No-XAI", "Dater", "Chain-of-Table", "Plan-of-SQLs"]

def save_session_data(username, data):
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        file_name = f'{username}_{timestamp}_session.json'
        json_data = json.dumps(data, indent=4)
        temp_file_path = f"/tmp/{file_name}"
        with open(temp_file_path, 'w') as f:
            f.write(json_data)
        api = HfApi()
        api.upload_file(
            path_or_fileobj=temp_file_path,
            path_in_repo=f"session_data_foward_simulation/{file_name}",
            repo_id="luulinh90s/Tabular-LLM-Study-Data",
            repo_type="space",
        )
        os.remove(temp_file_path)
        logger.info(f"Session data saved for user {username} in Hugging Face Data Space")
    except Exception as e:
        logger.exception(f"Error saving session data for user {username}: {e}")

def load_session_data(username):
    try:
        api = HfApi()
        files = api.list_repo_files(repo_id="luulinh90s/Tabular-LLM-Study-Data", repo_type="space")
        user_files = [f for f in files if f.startswith(f'session_data_foward_simulation/{username}_') and f.endswith('_session.json')]
        if not user_files:
            logger.warning(f"No session data found for user {username}")
            return None
        latest_file = sorted(user_files, reverse=True)[0]
        file_path = hf_hub_download(repo_id="luulinh90s/Tabular-LLM-Study-Data", repo_type="space", filename=latest_file)
        with open(file_path, 'r') as f:
            data = json.load(f)
        logger.info(f"Session data loaded for user {username} from Hugging Face Data Space")
        return data
    except Exception as e:
        logger.exception(f"Error loading session data for user {username}: {e}")
        return None

def load_samples():
    common_samples = []
    categories = ["TP", "TN", "FP", "FN"]

    for category in categories:
        files = set(os.listdir(f'htmls_NO_XAI/{category}'))
        for method in ["Dater", "Chain-of-Table", "Plan-of-SQLs"]:
            method_dir = VISUALIZATION_DIRS[method]
            files &= set(os.listdir(f'{method_dir}/{category}'))

        for file in files:
            common_samples.append({'category': category, 'file': file})

    logger.info(f"Found {len(common_samples)} common samples across all methods")
    return common_samples

def select_balanced_samples(samples):
    try:
        if len(samples) < 10:
            logger.warning(f"Not enough common samples. Only {len(samples)} available.")
            return samples

        selected_samples = random.sample(samples, 10)
        logger.info(f"Selected 10 unique samples")
        return selected_samples
    except Exception as e:
        logger.exception("Error selecting balanced samples")
        return []

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        username = request.form.get('username')
        seed = request.form.get('seed')
        method = request.form.get('method')
        if not username or not seed or not method:
            return "Please fill in all fields and select a method.", 400
        try:
            seed = int(seed)
            random.seed(seed)
            all_samples = load_samples()
            selected_samples = select_balanced_samples(all_samples)
            if len(selected_samples) == 0:
                return "No common samples were found", 500
            session_data = {
                'username': username,
                'seed': seed,
                'method': method,
                'selected_samples': selected_samples,
                'current_index': 0,
                'responses': [],
                'start_time': datetime.now().isoformat()
            }
            save_session_data(username, session_data)
            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:
            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))

        sample = selected_samples[current_index]
        visualization_dir = VISUALIZATION_DIRS[method]
        visualization_path = f"{visualization_dir}/{sample['category']}/{sample['file']}"

        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,
                               method=method)
    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'))

        session_data['responses'].append({
            'sample_id': session_data['current_index'],
            'user_prediction': prediction
        })

        session_data['current_index'] += 1
        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'))

        session_data['end_time'] = datetime.now().isoformat()
        responses = session_data['responses']
        method = session_data['method']

        json_file = f'Tabular_LLMs_human_study_vis_6_{method.upper().replace("-", "_")}.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]['file']
            index = visualization_file.split('-')[1].split('.')[0]

            ground_truth_key = f"{method.upper().replace('-', '_')}_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)

        session_data['accuracy'] = accuracy
        session_data['true_percentage'] = true_percentage
        session_data['false_percentage'] = false_percentage

        save_session_data(username, session_data)

        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}")
    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__":
    app.run(host="0.0.0.0", port=7860, debug=True)