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

app = Flask(__name__)

# Load model
model_id = "deepseek-ai/DeepSeek-V3-0324"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto")
model.to('cpu')

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

@app.route('/generate', methods=['POST'])
def generate():
    data = request.get_json()
    user_prompt = data['prompt']
    
    enhanced_prompt = (
        f"Generate a highly advanced, modern, single-file HTML website or web app with animations. "
        f"Use HTML, CSS, and JavaScript only. Return just the code. No explanation. Prompt:\n{user_prompt}\n\n```html\n"
    )

    inputs = tokenizer(enhanced_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=1000,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract HTML code only
    if "```html" in decoded:
        code = decoded.split("```html")[1].split("```")[0]
    else:
        code = decoded

    return jsonify({"code": code})


if __name__ == '__main__':
    app.run(debug=True)