import os, json, time
import requests
from flask import Flask, request, jsonify, Response, stream_with_context
from flask_cors import CORS

OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
app = Flask(__name__)
CORS(app)

def ollama_chat(messages, stream=False):
    return requests.post(
        f"{OLLAMA_URL}/api/chat",
        json={"model": "ornith:9b", "messages": messages, "stream": stream},
        stream=stream,
    )

@app.route("/")
def index():
    return jsonify({"service": "Ornith 9B via Ollama"})

@app.route("/health")
def health():
    try:
        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
        return jsonify({"status": "ok", "ollama": r.ok})
    except:
        return jsonify({"status": "error", "ollama": False})

@app.route("/ask", methods=["GET", "POST"])
def ask():
    stream = str(request.args.get("stream", "")).lower() in ["true", "1", "yes"]
    if request.method == "GET":
        q = request.args.get("q", "").strip()
    else:
        data = request.get_json(silent=True) or {}
        q = data.get("q", "").strip()
        if "stream" in data:
            stream = str(data.get("stream")).lower() in ["true", "1", "yes"]
    if not q:
        return jsonify({"error": "Missing q"}), 400

    r = ollama_chat([{"role": "user", "content": q}], stream=stream)

    if stream:
        def gen():
            for line in r.iter_lines(decode_unicode=True):
                if line:
                    try:
                        d = json.loads(line)
                        if d.get("done"):
                            yield "data: [DONE]\n\n"
                        else:
                            yield f"data: {json.dumps({'text': d.get('message', {}).get('content', '')})}\n\n"
                    except:
                        pass
        return Response(gen(), mimetype="text/event-stream")
    else:
        data = r.json()
        return jsonify({"response": data.get("message", {}).get("content", "")})

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    print(f"[INFO] Server on :{port}", flush=True)
    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)