from flask import Flask, request, jsonify
from sentence_transformers import SentenceTransformer
import logging
import os
from functools import wraps

logging.basicConfig(level=logging.DEBUG)

API_TOKEN = os.getenv("API_TOKEN", "your-default-token")  # Set a default for local testing


app = Flask(__name__)

def require_token(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth_header = request.headers.get("Authorization")
        if not auth_header:
            return jsonify({"error": "Authorization header missing"}), 401
        try:
            token = auth_header.split("Bearer ")[1]
            if token != API_TOKEN:
                return jsonify({"error": "Invalid token"}), 403
        except IndexError:
            return jsonify({"error": "Invalid Authorization header format"}), 401
        return f(*args, **kwargs)
    return decorated
    

# Load the Qwen3-Embedding-0.6B model (adjust the path as needed)
model = SentenceTransformer(model_name_or_path="Qwen/Qwen3-Embedding-0.6B")

@app.route('/')
def home():
    return 'Hello from Flask on Hugging Face!'
    
@app.route('/embed', methods=['POST'])
@require_token
def get_embedding():
    try:
        text = request.json.get('text', '')
        if not text:
            return jsonify({"error": "No text provided"}), 400
        
        # Generate embeddings for the input text
        document_embeddings = model.encode(text)
        arr_list = document_embeddings.tolist()
        return jsonify({"embedding": arr_list})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


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