from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
from transformers import T5ForConditionalGeneration, T5Tokenizer
import torch
import logging
import os

app = Flask(__name__)
CORS(app)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load model from Hugging Face Hub
try:
    logger.info("Loading model from Hugging Face Hub...")
    model_id = "Mayank14/t5-simplifier-model"
    tokenizer = T5Tokenizer.from_pretrained(model_id)
    model = T5ForConditionalGeneration.from_pretrained(model_id)
    logger.info("Model loaded successfully!")
except Exception as e:
    logger.error(f"Failed to load model: {e}")
    raise

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

@app.route('/simplify', methods=['POST'])
def simplify():
    try:
        data = request.get_json()
        if not data or 'text' not in data:
            return jsonify({'error': 'No text provided'}), 400

        text = data['text'].strip()
        if not text:
            return jsonify({'error': 'Empty text provided'}), 400

        if len(text) > 2000:
            return jsonify({'error': 'Text too long. Please limit to 2000 characters.'}), 400

        # Prepare input
        input_text = "simplify: " + text
        inputs = tokenizer.encode(
            input_text,
            return_tensors="pt",
            max_length=512,
            truncation=True,
            padding=True
        )

        # Generate simplified text
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=512,
                num_beams=4,
                early_stopping=True,
                do_sample=False,
                temperature=0.7
            )

        simplified_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Clean up output
        if simplified_text.lower().startswith("simplify:"):
            simplified_text = simplified_text[9:].strip()

        return jsonify({
            'simplified_text': simplified_text,
            'original_length': len(text),
            'simplified_length': len(simplified_text)
        })

    except Exception as e:
        logger.error(f"Error in simplify endpoint: {e}")
        return jsonify({'error': 'Internal server error'}), 500

@app.route('/health')
def health():
    return jsonify({
        'status': 'healthy',
        'model_loaded': True
    })

if __name__ == '__main__':
    # Get port from environment variable for Hugging Face Spaces
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port)