# ==========================================================
# 🚀 NEXA AI - STABLE OLLAMA COLAB API SERVER (FIXED)
# ==========================================================

import os
import time
import threading
import subprocess
import requests
import json
from flask import Flask, request, jsonify, Response
from flask_cors import CORS
from pyngrok import ngrok

# ---------------- CONFIG ----------------
MODEL_NAME = "phi3:mini"   # fast model
PORT = 5000
NGROK_AUTH_TOKEN = ""      # Task 37: Show warning if empty

# ---------------- INSTALL ----------------
print("📦 Installing dependencies...")

subprocess.run("apt-get update -y", shell=True)
subprocess.run("apt-get install -y zstd", shell=True)

subprocess.run("curl -fsSL https://ollama.com/install.sh | sh", shell=True)
subprocess.run("pip install flask flask-cors pyngrok requests", shell=True)

# ---------------- START OLLAMA ----------------
print("🧠 Starting Ollama...")

def start_ollama():
    os.environ["OLLAMA_HOST"] = "0.0.0.0:11434"
    subprocess.run("ollama serve", shell=True)

ollama_thread = threading.Thread(target=start_ollama, daemon=True)
ollama_thread.start()

# ---------------- WAIT FOR OLLAMA ----------------
print("⏳ Waiting for Ollama to be ready...")

def wait_for_ollama():
    while True:
        try:
            r = requests.get("http://localhost:11434")
            if r.status_code == 200:
                break
        except:
            pass
        time.sleep(2)

wait_for_ollama()
print("✅ Ollama is ready!")

# ---------------- PULL MODEL ----------------
print(f"📥 Pulling model: {MODEL_NAME}")
subprocess.run(f"ollama pull {MODEL_NAME}", shell=True)

# ---------------- FLASK API ----------------
app = Flask(__name__)
CORS(app)

@app.route("/status", methods=["GET"])
def status():
    return jsonify({
        "status": "online",
        "model": MODEL_NAME,
        "backend": "ollama"
    })

# Task 35 & 36: Implement /api/chat with streaming and dynamic system prompt
@app.route("/api/chat", methods=["POST"])
def chat():
    data = request.json or {}
    messages = data.get("messages", [])
    model = data.get("model", MODEL_NAME)
    stream = data.get("stream", True)
    
    if not messages:
        return jsonify({"error": "No messages provided"}), 400

    def generate():
        try:
            # Task 35: Use Ollama's chat API directly
            response = requests.post(
                "http://localhost:11434/api/chat",
                json={
                    "model": model,
                    "messages": messages,
                    "stream": stream
                },
                stream=True,
                timeout=120
            )

            for line in response.iter_lines():
                if line:
                    chunk = json.loads(line)
                    # Task 35: Forward the chunk as NDJSON
                    yield json.dumps(chunk) + "\n"

        except Exception as e:
            yield json.dumps({"error": str(e), "message": {"content": "Colab server error."}}) + "\n"

    if stream:
        return Response(generate(), mimetype='application/x-ndjson')
    else:
        try:
            response = requests.post(
                "http://localhost:11434/api/chat",
                json={
                    "model": model,
                    "messages": messages,
                    "stream": False
                },
                timeout=120
            )
            return jsonify(response.json())
        except Exception as e:
            return jsonify({"error": str(e)}), 500

def run_api():
    app.run(host="0.0.0.0", port=PORT)

threading.Thread(target=run_api, daemon=True).start()

# ---------------- NGROK ----------------
print("🔗 Starting public tunnel...")

# Task 37: ngrok auth token validation
if not NGROK_AUTH_TOKEN:
    print("\n" + "!"*50)
    print("⚠️ WARNING: NGROK_AUTH_TOKEN is empty!")
    print("Tunnels will be limited and expire quickly.")
    print("Get your token at: https://dashboard.ngrok.com/get-started/your-authtoken")
    print("!"*50 + "\n")
else:
    ngrok.set_auth_token(NGROK_AUTH_TOKEN)

try:
    public_url = ngrok.connect(PORT)
    print("\n" + "="*50)
    print("🚀 NEXA AI IS LIVE")
    print("📡 URL:", public_url.public_url)
    print("🤖 Model:", MODEL_NAME)
    print("="*50)
except Exception as e:
    print(f"❌ Failed to start ngrok: {e}")

# ---------------- KEEP ALIVE ----------------
while True:
    try:
        time.sleep(60)
    except KeyboardInterrupt:
        print("Stopped")
        break