from flask import Flask, request, Response, jsonify
import requests
import threading
import time
import random
import logging
from typing import Optional

app = Flask(__name__)

# =========================
# CONFIG
# =========================
Server  = [
    "https://Techbitforge-hello_server.hf.space"
]
HF_SERVERS = [
    "https://techbitforge-ollamaserver.hf.space/"
]

PING_INTERVAL = 7200
REQUEST_TIMEOUT = 300

# =========================
# LOGGING
# =========================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s — %(levelname)s — %(message)s"
)
logger = logging.getLogger("proxy")


def pick_server() -> str:
    return random.choice(HF_SERVERS)


def error_response(message: str, code: int = 500):
    logger.error(message)
    return jsonify({
        "error": {
            "message": message,
            "type": "proxy_error"
        }
    }), code


# =========================
# CHAT COMPLETIONS PROXY
# =========================
@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    try:
        payload = request.get_json(force=True)
    except Exception:
        return error_response("Invalid JSON payload", 400)

    headers = dict(request.headers)
    headers.pop("Host", None)

    target = pick_server() + "/v1/chat/completions"
    stream = payload.get("stream", False)

    logger.info(f"Forwarding request to {target} | Stream={stream}")

    try:
        if stream:
            def generate():
                with requests.post(
                    target,
                    json=payload,
                    headers=headers,
                    stream=True,
                    timeout=REQUEST_TIMEOUT
                ) as res:
                    res.raise_for_status()
                    for line in res.iter_lines():
                        if line:
                            yield line + b"\n\n"

            return Response(
                generate(),
                mimetype="text/event-stream",
                headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
            )

        res = requests.post(
            target,
            json=payload,
            headers=headers,
            timeout=REQUEST_TIMEOUT
        )
        res.raise_for_status()
        return jsonify(res.json())

    except requests.exceptions.Timeout:
        return error_response("Upstream server timeout", 504)
    except requests.exceptions.ConnectionError:
        return error_response("Failed to connect to upstream server", 502)
    except Exception as e:
        return error_response(str(e), 500)


# =========================
# MODELS ROUTE
# =========================
@app.route("/models", methods=["GET"])
def models():
    url = pick_server()
    target = f"{url}/models"

    logger.info(f"Fetching models from {target}")

    try:
        res = requests.get(target, timeout=20)
        res.raise_for_status()
        return jsonify(res.json())
    except Exception as e:
        return error_response(f"Failed to fetch models: {e}", 500)


# =========================
# HEALTH CHECK
# =========================
@app.route("/", methods=["GET"])
def home():
    return jsonify({"status": "running", "service": "HF Proxy"})


import requests
import threading
import time


PING_INTERVAL = 3600  # 24 hours


def keep_alive_worker():
    url = "https://techbitforge-server.hf.space/"

    while True:
        try:
            response = requests.get(url, timeout=120)

            if response.status_code == 200:
                print("Success: server is alive")
            else:
                print(f"Server returned status code: {response.status_code}")

        except requests.RequestException as e:
            print(f"Ping failed: {e}")

        # Wait 24 hours before the next ping
        time.sleep(PING_INTERVAL)


def start_background_thread():
    thread = threading.Thread(
        target=keep_alive_worker,
        daemon=True
    )

    thread.start()

# =========================
# MAIN
# =========================
if __name__ == "__main__":
    start_background_thread()
    app.run(
        host="0.0.0.0",
        port=7860,
        threaded=True,
        debug=False
    )