# demo_server.py — run on H100, connects to Flutter app via WiFi
from flask import Flask, request, jsonify
import torch
import logging

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

# Note: Loading the 15GB model takes significant VRAM. 
# For production, we load this once during server startup.
model = None
tokenizer = None
adapters = {}

def load_model():
    global model, tokenizer
    if model is None:
        log.info("Loading Base Model (15GB)...")
        from transformers import AutoModelForCausalLM, AutoTokenizer
        
        # Load the base Ankahi merged model
        # We use bfloat16 to match training dtype
        model = AutoModelForCausalLM.from_pretrained(
            "punjabi_gemma/ankahi",
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
        tokenizer = AutoTokenizer.from_pretrained("punjabi_gemma/ankahi")
        
        # Load the 5 persona adapters into the PEFT model
        from peft import PeftModel
        log.info("Loading Persona Adapters...")
        model = PeftModel.from_pretrained(model, "artifacts/stage2/ananya", adapter_name="ananya")
        model.load_adapter("artifacts/stage2/arjun", adapter_name="arjun")
        model.load_adapter("artifacts/stage2/priya", adapter_name="priya")
        model.load_adapter("artifacts/stage2/rohan", adapter_name="rohan")
        model.load_adapter("artifacts/stage2/zara", adapter_name="zara")
        
        model.eval()
        log.info("Server Ready!")

@app.route("/predict", methods=["POST"])
def predict():
    if model is None:
        return jsonify({"error": "Model not loaded yet. Call /init first or wait."}), 503
        
    data = request.json
    if not data or "pictograms" not in data or "persona" not in data:
        return jsonify({"error": "Invalid request. Need 'pictograms' array and 'persona' string."}), 400
        
    persona = data["persona"]
    pictograms = data["pictograms"]
    context = data.get("context", "")
    
    # Activate the specific child's adapter
    try:
        model.set_adapter(persona)
    except Exception as e:
        return jsonify({"error": f"Unknown persona '{persona}'. {e}"}), 400
        
    # Format the prompt
    system_prompt = "Translate the pictograms into a natural, spoken sentence."
    user_msg = f"Context: {context}\n" if context else ""
    user_msg += f"Pictograms: {', '.join(pictograms)}"
    
    prompt = f"<bos><start_of_turn>user\n{system_prompt}\n{user_msg}<end_of_turn>\n<start_of_turn>model\n"
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=30, 
            num_return_sequences=3, # Return 3 alternatives
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        
    alternatives = []
    for output in outputs:
        # Decode only the generated part
        generated_text = tokenizer.decode(output[inputs.input_ids.shape[1]:], skip_special_tokens=True)
        alternatives.append(generated_text.strip())
        
    return jsonify({
        "persona": persona,
        "pictograms": pictograms,
        "alternatives": alternatives
    })

@app.route("/init", methods=["GET"])
def init():
    load_model()
    return jsonify({"status": "ready"})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)