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

app = Flask(__name__)

ENDPOINT_URL = "https://o1zwshreete15x04.us-east-1.aws.endpoints.huggingface.cloud"
HF_TOKEN     = os.getenv("HF_TOKEN", "")
CHUNK_SIZE   = int(os.getenv("CHUNK_SIZE", "2"))    # chars per stream chunk
CHUNK_DELAY  = float(os.getenv("CHUNK_DELAY", "0.012"))  # seconds between chunks


def _headers():
    return {
        "Accept":        "application/json",
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type":  "application/json",
    }


@app.route("/health", methods=["GET"])
def health():
    """
    Lightweight health check — just pings the endpoint with a tiny prompt.
    """
    if not ENDPOINT_URL or not HF_TOKEN:
        return jsonify({
            "status":  "unhealthy",
            "error":   "ENDPOINT_URL or HF_TOKEN env var not set"
        }), 503

    try:
        resp = requests.post(
            ENDPOINT_URL,
            headers=_headers(),
            json={
                "inputs": [{"role": "user", "content": "Hi"}],
                "parameters": {
                    "max_new_tokens": 5,
                    "temperature":    0.1,
                    "do_sample":      False,
                }
            },
            timeout=20,
        )
        if resp.status_code == 200:
            return jsonify({"status": "healthy", "code": 200})
        return jsonify({
            "status": "unhealthy",
            "code":   resp.status_code,
            "error":  resp.text[:300],
        }), 503

    except requests.exceptions.Timeout:
        return jsonify({"status": "unhealthy", "error": "Endpoint timed out"}), 503
    except Exception as e:
        return jsonify({"status": "unhealthy", "error": str(e)}), 503


@app.route("/chat", methods=["POST"])
def chat():
    """
    Accepts:
    {
        "messages":    [{role, content}, ...],
        "max_tokens":  512,
        "temperature": 0.7,
        "top_p":       0.9,
        "do_sample":   true,
        "stream":      true      <-- if true, SSE stream back to caller
    }
    """
    data = request.get_json(silent=True) or {}

    messages    = data.get("messages", [])
    max_tokens  = int(data.get("max_tokens", 512))
    temperature = float(data.get("temperature", 0.7))
    top_p       = float(data.get("top_p", 0.9))
    do_sample   = bool(data.get("do_sample", temperature > 0))
    stream      = bool(data.get("stream", True))

    if not messages:
        return jsonify({"error": "messages array required"}), 400

    # Build payload for your custom handler
    payload = {
        "inputs": [{"role": m["role"], "content": m["content"]} for m in messages],
        "parameters": {
            "max_new_tokens": max_tokens,
            "temperature":    temperature,
            "top_p":          top_p,
            "do_sample":      do_sample,
        }
    }

    # ── Non-streaming ─────────────────────────────────────────
    if not stream:
        try:
            resp = requests.post(
                ENDPOINT_URL,
                headers=_headers(),
                json=payload,
                timeout=90,
            )
            resp.raise_for_status()
            result = resp.json()
            text = _extract_text(result)
            return jsonify({"generated_text": text, "ok": True})
        except Exception as e:
            return jsonify({"error": str(e)}), 500

    # ── Streaming ─────────────────────────────────────────────
    def generate():
        try:
            resp = requests.post(
                ENDPOINT_URL,
                headers=_headers(),
                json=payload,
                timeout=90,
            )
            resp.raise_for_status()
            result    = resp.json()
            full_text = _extract_text(result)

            if not full_text:
                yield _sse({"error": "Empty response from model"})
                return

            # ── Smooth streaming in small character chunks ────
            # We buffer into "word-aware" chunks so words don't
            # get split mid-character in a jarring way.
            buffer = ""
            for char in full_text:
                buffer += char

                # Flush on punctuation/spaces for natural rhythm
                should_flush = (
                    len(buffer) >= CHUNK_SIZE or
                    char in (' ', '\n', '.', ',', '!', '?', ':', ';')
                )

                if should_flush and buffer:
                    yield _sse({"token": buffer})
                    buffer = ""
                    time.sleep(CHUNK_DELAY)

            # Flush any remaining buffer
            if buffer:
                yield _sse({"token": buffer})

            yield "data: [DONE]\n\n"

        except requests.exceptions.Timeout:
            yield _sse({"error": "Endpoint timed out — try again"})
        except requests.exceptions.HTTPError as e:
            yield _sse({"error": f"Endpoint error {e.response.status_code}: {e.response.text[:200]}"})
        except Exception as e:
            yield _sse({"error": str(e)})

    return Response(
        stream_with_context(generate()),
        content_type="text/event-stream",
        headers={
            "Cache-Control":     "no-cache",
            "X-Accel-Buffering": "no",
            "Connection":        "keep-alive",
            "Access-Control-Allow-Origin": "*",
        }
    )


@app.route("/", methods=["GET"])
def root():
    return jsonify({
        "name":    "SmilyAI Proxy",
        "status":  "running",
        "routes":  ["/health", "/chat"],
        "model":   "SmilyAI 1.2B ChatML",
    })


def _extract_text(result):
    """Pull generated_text out of whatever shape the endpoint returns."""
    if isinstance(result, list) and len(result) > 0:
        return result[0].get("generated_text", "")
    if isinstance(result, dict):
        return result.get("generated_text", "")
    return str(result)


def _sse(obj):
    """Format a dict as an SSE data line."""
    return f"data: {json.dumps(obj)}\n\n"


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)