from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import os

app = Flask(__name__)

# Dùng mô hình nhỏ hơn để tiết kiệm bộ nhớ
model_name = "VietAI/vit5-base-vietnews-summarization"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

summarizer = pipeline(
    "summarization",
    model=model,
    tokenizer=tokenizer
)

@app.route("/")
def home():
    return jsonify({
        "message": "✅ AI Text Summarization API is running (Vietnamese).",
        "usage": "POST /summarize với JSON: { 'text': '...', 'topic': '...' }"
    })

@app.route("/summarize", methods=["POST"])
def summarize():
    data = request.get_json(force=True)

    if not data or "text" not in data:
        return jsonify({"error": "Thiếu trường 'text'"}), 400

    text = data["text"]
    topic = data.get("topic", "").strip()  # Nếu không có topic, mặc định trống

    # Nếu có topic, ghép vào đầu văn bản để mô hình hiểu hướng tóm tắt
    if topic:
        text = f"Chủ đề: {topic}. Nội dung: {text}"

    # Giới hạn số từ để tránh lỗi bộ nhớ
    max_words = 1000
    if len(text.split()) > max_words:
        text = " ".join(text.split()[:max_words])

    try:
        # Tóm tắt văn bản
        result = summarizer(text, max_length=150, min_length=40, do_sample=False)
        summary = result[0]['summary_text']

        return jsonify({"summary": summary})

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port)