import os
import json
import requests as _requests
from bs4 import BeautifulSoup
from flask import Flask, render_template, request, jsonify, stream_with_context, Response
from airllm import AutoModel

app = Flask(__name__)

model = None
chat_history = []

MODEL_NAME = "Exocore/ExocoreV1"


def load_model():
    global model
    try:
        model = AutoModel.from_pretrained(MODEL_NAME)
    except Exception as e:
        print(f"Model load error: {e}")


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


@app.route("/api/model-details", methods=["GET"])
def model_details():
    if model is None:
        return jsonify({"error": "No model loaded"}), 400
    return jsonify({
        "name": "ExocoreV1",
        "type": "airllm",
    })


@app.route("/api/chat", methods=["POST"])
def chat():
    global chat_history
    if model is None:
        return jsonify({"error": "Model not loaded"}), 400

    data = request.json
    user_message = data.get("message", "").strip()
    max_tokens = int(data.get("max_tokens", 128))

    if not user_message:
        return jsonify({"error": "Empty message"}), 400

    def generate():
        full_reply = ""
        try:
            input_text = [user_message]
            input_tokens = model.tokenizer(
                input_text,
                return_tensors="pt",
                return_attention_mask=False,
                truncation=True,
                max_length=128,
                padding=False
            )
            generation_output = model.generate(
                input_tokens["input_ids"].cuda(),
                max_new_tokens=max_tokens,
                use_cache=True,
                return_dict_in_generate=True
            )
            output = model.tokenizer.decode(generation_output.sequences[0])
            full_reply = output
            for ch in full_reply:
                yield f"data: {json.dumps({'token': ch})}\n\n"
            chat_history.append({"role": "user", "content": user_message})
            chat_history.append({"role": "assistant", "content": full_reply})
            yield f"data: {json.dumps({'done': True})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return Response(stream_with_context(generate()), mimetype="text/event-stream")


@app.route("/api/clear", methods=["POST"])
def clear_history():
    global chat_history
    chat_history = []
    return jsonify({"success": True})


@app.route("/api/status", methods=["GET"])
def status():
    return jsonify({
        "model_loaded": model is not None,
        "history_length": len(chat_history),
    })


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