#!/usr/bin/env python3
"""
SmilyAI Proxy — runs on a HF Space.
Forwards authenticated requests to a private HF Inference Endpoint.
Exists because PythonAnywhere's free tier can't reach arbitrary outbound
hosts, but *.hf.space IS whitelisted.
"""

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

app = Flask(__name__)
CORS(app)

# ── Config (set these as SECRETS in the Space settings, not hardcoded!) ──────
TARGET_BASE_URL = os.environ.get("TARGET_BASE_URL", "").rstrip("/")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
PROXY_SECRET = os.environ.get("PROXY_SECRET", "")  # shared secret w/ your Flask app

if not TARGET_BASE_URL or not HF_TOKEN or not PROXY_SECRET:
    print("⚠ WARNING: TARGET_BASE_URL / HF_TOKEN / PROXY_SECRET not fully set!")


def check_auth(req) -> bool:
    key = req.headers.get("X-Proxy-Key", "")
    return PROXY_SECRET and key == PROXY_SECRET


@app.route("/", methods=["GET"])
def health():
    return jsonify({"status": "ok", "service": "smilyai-proxy"})


@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    if not check_auth(request):
        return jsonify({"error": "unauthorized"}), 401

    try:
        payload = request.get_json(force=True)
    except Exception:
        return jsonify({"error": "invalid json body"}), 400

    is_stream = bool(payload.get("stream", False))
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json; charset=utf-8",  # ← add charset here
    }
    url = f"{TARGET_BASE_URL}/chat/completions"

    if not is_stream:
        try:
            r = requests.post(url, headers=headers, json=payload, timeout=120)
            return Response(
                r.content,
                status=r.status_code,
                content_type="application/json; charset=utf-8"  # ← fix here too
            )
        except Exception as e:
            return jsonify({"error": str(e)}), 502

    def generate():
        try:
            with requests.post(
                url, headers=headers, json=payload,
                stream=True, timeout=120
            ) as r:
                # ← KEY FIX: iter_lines with encoding="utf-8" explicitly
                for raw_line in r.iter_lines(decode_unicode=False):
                    if raw_line is None:
                        continue
                    # decode bytes as utf-8 explicitly
                    if isinstance(raw_line, bytes):
                        line = raw_line.decode("utf-8", errors="replace")
                    else:
                        line = raw_line
                    if not line:
                        continue
                    yield (line + "\n\n").encode("utf-8")
        except Exception as e:
            err = json.dumps({"error": str(e)}, ensure_ascii=False)
            yield f"data: {err}\n\n".encode("utf-8")
            yield b"data: [DONE]\n\n"

    return Response(
        generate(),
        content_type="text/event-stream; charset=utf-8",  # ← charset here
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Transfer-Encoding": "chunked",
        }
    )

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