import os
import pickle
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences
from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
from huggingface_hub import hf_hub_download

app = Flask(__name__)
CORS(app)

# Configuration
MAX_LEN = 300
REPO_ID = "d-e-e-k-11/chatbot-performance-analyzer"

# Global variables
model = None
tokenizer = None

# Custom Attention Layer
class Attention(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(Attention, self).__init__(**kwargs)
    def build(self, input_shape):
        self.W = self.add_weight(name='attention_weight', shape=(input_shape[-1], 1), initializer='random_normal', trainable=True)
        self.b = self.add_weight(name='attention_bias', shape=(input_shape[1], 1), initializer='zeros', trainable=True)
        super(Attention, self).build(input_shape)
    def call(self, x):
        e = tf.keras.backend.tanh(tf.keras.backend.dot(x, self.W) + self.b)
        a = tf.keras.backend.softmax(e, axis=1)
        output = x * a
        return tf.keras.backend.sum(output, axis=1)

def load_resources():
    global model, tokenizer
    try:
        # Download from Hugging Face Hub
        model_path = hf_hub_download(repo_id=REPO_ID, filename="chatbot_performance_advanced.h5", repo_type="space")
        tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer_advanced.pickle", repo_type="space")
        
        model = tf.keras.models.load_model(model_path, custom_objects={'Attention': Attention})
        with open(tokenizer_path, 'rb') as handle:
            tokenizer = pickle.load(handle)
        print("Advanced Model and Tokenizer loaded successfully from Hub.")
        return True
    except Exception as e:
        print(f"Error loading resources: {e}")
        return False

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    global model, tokenizer
    if model is None or tokenizer is None:
        if not load_resources():
            return jsonify({
                'error': 'Model is still loading. Please wait a moment.',
                'status': 'loading'
            }), 503

    data = request.json
    if data.get('ping'):
        return jsonify({'status': 'ready'})

    facts = data.get('facts', 'No context provided.')
    question = data.get('question', '')
    response = data.get('response', '')
    
    try:
        # Preprocess text with facts context
        text = f"[FACTS] {facts} [QUERY] {question} [RES] {response}".lower()
        seq = tokenizer.texts_to_sequences([text])
        pad = pad_sequences(seq, maxlen=MAX_LEN)
        
        # Prediction
        prediction = model.predict(pad)[0][0]
        is_best = bool(prediction > 0.5)
        
        return jsonify({
            'probability': float(prediction),
            'is_best': is_best
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

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