main.py

from flask import Flask, request, jsonify, send_from_directory
from model_handler import ModelHandler
import supabase
import os
import random
from datetime import datetime, timedelta
import smtplib
from email.mime.text import MIMEText
import jwt

# --------------------------
# Config
# --------------------------
SUPABASE_URL = os.environ["SUPABASE_URL"]
SUPABASE_KEY = os.environ["SUPABASE_KEY"]
JWT_SECRET = os.environ.get("JWT_SECRET", "supersecretkey")
JWT_ALGORITHM = "HS256"
EMAIL_HOST = os.environ.get("EMAIL_HOST", "smtp.gmail.com")
EMAIL_PORT = int(os.environ.get("EMAIL_PORT", 587))
EMAIL_USER = os.environ.get("EMAIL_USER")
EMAIL_PASS = os.environ.get("EMAIL_PASS")

# --------------------------
# Initialize Supabase
# --------------------------
sb = supabase.create_client(SUPABASE_URL, SUPABASE_KEY)

# --------------------------
# Initialize Flask
# --------------------------
app = Flask(__name__, static_folder=".", template_folder=".")

# --------------------------
# Models Initialization
# --------------------------
loaded_models = {
    "qwen": ModelHandler("Qwen/Qwen2.5-Coder-1.5B"),
    "deepseek": ModelHandler("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
}
dynamic_model_cache = {}

# --------------------------
# Helpers
# --------------------------
def generate_verification_code():
    return str(random.randint(100000, 999999))

def send_verification_email(to_email, code):
    try:
        msg = MIMEText(f"Your verification code is: {code}")
        msg['Subject'] = "AI_HUB Email Verification"
        msg['From'] = EMAIL_USER
        msg['To'] = to_email

        server = smtplib.SMTP(EMAIL_HOST, EMAIL_PORT)
        server.starttls()
        server.login(EMAIL_USER, EMAIL_PASS)
        server.send_message(msg)
        server.quit()
    except Exception as e:
        print("Error sending email:", e)

def create_jwt(user_id, email, username):
    payload = {
        "user_id": user_id,
        "email": email,
        "username": username,
        "exp": datetime.utcnow() + timedelta(days=30)
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)

def decode_jwt(token):
    try:
        return jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

def auth_required(f):
    def wrapper(*args, **kwargs):
        token = request.headers.get("Authorization")
        if not token:
            return jsonify({"error": "Unauthorized"}), 401
        if token.startswith("Bearer "):
            token = token.split(" ")[1]
        user_data = decode_jwt(token)
        if not user_data:
            return jsonify({"error": "Invalid or expired token"}), 401
        request.user = user_data
        return f(*args, **kwargs)
    wrapper.__name__ = f.__name__
    return wrapper

# --------------------------
# User Endpoints
# --------------------------
@app.route("/api/signup", methods=["POST"])
def signup():
    data = request.get_json()
    email = data.get("email")
    username = data.get("username")
    password = data.get("password")

    existing = sb.table("users").select("*").or_(f"email.eq.{email},username.eq.{username}").execute()
    if existing.data:
        return jsonify({"error": "Email or username already exists"}), 400

    code = generate_verification_code()
    expiry = datetime.utcnow() + timedelta(minutes=10)

    sb.table("users").insert({
        "email": email,
        "username": username,
        "password": password,
        "email_verified": False,
        "email_verification_code": code,
        "code_expires_at": expiry
    }).execute()

    send_verification_email(email, code)

    return jsonify({"message": "Signup successful. Verify your email.", "email_verification_code": code})

@app.route("/api/verify_email", methods=["POST"])
def verify_email():
    data = request.get_json()
    email = data.get("email")
    code = data.get("code")

    res = sb.table("users").select("*").eq("email", email).execute()
    if not res.data:
        return jsonify({"error": "User not found"}), 404

    user = res.data[0]
    if user["email_verification_code"] != code:
        return jsonify({"error": "Incorrect verification code"}), 400
    if user["code_expires_at"] < datetime.utcnow():
        return jsonify({"error": "Verification code expired"}), 400

    sb.table("users").update({
        "email_verified": True,
        "email_verification_code": None,
        "code_expires_at": None
    }).eq("email", email).execute()

    return jsonify({"message": "Email verified successfully"})

@app.route("/api/login", methods=["POST"])
def login():
    data = request.get_json()
    email = data.get("email")
    password = data.get("password")

    res = sb.table("users").select("*").eq("email", email).execute()
    users = res.data
    if not users:
        return jsonify({"error": "User not found"}), 404

    user = users[0]
    if user["password"] != password:
        return jsonify({"error": "Incorrect password"}), 401

    token = create_jwt(user["id"], user["email"], user["username"])
    return jsonify({
        "message": "Login successful",
        "token": token,
        "user": {"email": user["email"], "username": user["username"]}
    })

# --------------------------
# Session Endpoint
# --------------------------
@app.route("/api/session", methods=["GET"])
def session_check():
    token = request.headers.get("Authorization")
    if not token:
        return jsonify({"error": "No token provided"}), 401

    if token.startswith("Bearer "):
        token = token.split(" ")[1]

    user_data = decode_jwt(token)
    if not user_data:
        return jsonify({"error": "Invalid or expired token"}), 401

    return jsonify({
        "user": {
            "user_id": user_data["user_id"],
            "email": user_data["email"],
            "username": user_data["username"]
        }
    })

# --------------------------
# Profile Endpoint
# --------------------------
@app.route("/api/profile/<uuid:user_id>", methods=["GET"])
@auth_required
def get_profile(user_id):
    res = sb.table("profiles").select("*").eq("user_id", str(user_id)).execute()
    if not res.data:
        return jsonify({"error": "Profile not found"}), 404
    return jsonify(res.data[0])

# --------------------------
# Chat Endpoint
# --------------------------
@app.route("/api/chat", methods=["POST"])
@auth_required
def chat():
    data = request.get_json()
    message = data.get("message", "")
    model_choice = data.get("model", "qwen")
    max_tokens = data.get("max_tokens", 200)

    model = loaded_models.get(model_choice)
    if model_choice.startswith("http"):
        if model_choice in dynamic_model_cache:
            model = dynamic_model_cache[model_choice]
        else:
            try:
                model = ModelHandler(model_choice)
                dynamic_model_cache[model_choice] = model
            except Exception as e:
                return jsonify({"error": f"Failed to load model: {e}"}), 500

    if not model:
        return jsonify({"error": "Invalid model"}), 400

    try:
        response = model.chat(message, max_new_tokens=max_tokens)
    except Exception as e:
        response = f"Error generating response: {str(e)}"

    return jsonify({"response": response})

# --------------------------
# AI Apps Endpoints
# --------------------------
@app.route("/api/apps", methods=["GET"])
def list_apps():
    res = sb.table("bots").select("*").execute()
    return jsonify(res.data)

@app.route("/api/apps/create", methods=["POST"])
@auth_required
def create_app():
    data = request.get_json()
    name = data.get("name")
    language = data.get("language")
    code = data.get("code")
    api_key = generate_verification_code()

    sb.table("bots").insert({
        "name": name,
        "language": language,
        "code": code,
        "creator_id": request.user["user_id"],
        "api_key": api_key
    }).execute()

    return jsonify({"message": "AI app created successfully", "api_key": api_key})

# --------------------------
# Comments Endpoint
# --------------------------
@app.route("/api/comments/<int:app_id>", methods=["GET"])
def get_comments(app_id):
    res = sb.table("comments").select("*").eq("server_id", app_id).execute()
    return jsonify(res.data)

@app.route("/api/comments/<int:app_id>", methods=["POST"])
@auth_required
def post_comment(app_id):
    data = request.get_json()
    content = data.get("content")
    sb.table("comments").insert({
        "server_id": app_id,
        "user_id": request.user["user_id"],
        "content": content
    }).execute()
    return jsonify({"message": "Comment posted"})

# --------------------------
# Serve frontend
# --------------------------
@app.route("/")
def index():
    return send_from_directory(".", "index.html")

@app.route("/<path:path>")
def static_files(path):
    return send_from_directory(".", path)

# --------------------------
# Run App
# --------------------------
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)