"""
GPT-2 Text Generation Server
=============================
A minimal Flask API that loads GPT-2 from Hugging Face Transformers and exposes
a single POST /generate endpoint. The model is loaded once at startup and kept
in memory, so generation requests are fast after the first cold start (~10-30s).

Endpoints
---------
POST /generate
    Body:  { "prompt": str, "max_new_tokens": int, "temperature": float,
             "top_p": float, "repetition_penalty": float }
    Returns: { "generated_text": str, "prompt": str }

GET /
    Serves the frontend HTML page.

Usage
-----
    pip install flask transformers torch
    python app.py
    # Open http://localhost:5000
"""

from flask import Flask, request, jsonify, send_from_directory
from transformers import GPT2LMHeadModel, AutoTokenizer
import torch
import os

app = Flask(__name__, static_folder="static")

print("Loading GPT-2 model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("helloadhavan/llara1.1-100M-base")
model = GPT2LMHeadModel.from_pretrained("helloadhavan/llara1.1-100M-base")
model.eval()
print("Model ready.")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model.to(DEVICE)


@app.route("/")
def index():
    return send_from_directory(".", "index.html")


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

    prompt = body.get("prompt", "").strip()
    if not prompt:
        return jsonify({"error": "prompt is required"}), 400

    max_new_tokens = int(body.get("max_new_tokens", 200))
    temperature = float(body.get("temperature", 0.9))
    top_p = float(body.get("top_p", 0.95))
    repetition_penalty = float(body.get("repetition_penalty", 1.2))

    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            pad_token_id=tokenizer.eos_token_id,
        )

    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    generated_only = full_text[len(prompt):]

    return jsonify({"generated_text": generated_only, "prompt": prompt})


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=7860)