import os
import subprocess
import time
import requests
import threading
from huggingface_hub import hf_hub_download
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
import uvicorn

APP_DIR = "/app"
LLAMA_DIR = os.path.join(APP_DIR, "llama.cpp")
BIN = os.path.join(LLAMA_DIR, "build", "bin", "llama-server")
MODEL_DIR = os.path.join(APP_DIR, "models")
LLAMA_PORT = 8080

os.makedirs(MODEL_DIR, exist_ok=True)

# متغير عام لتتبع حالة الخادم الداخلي
backend_ready = False


def build_llama_cpp():
    """بناء llama.cpp من المصدر إن لم يكن الملف التنفيذي موجودًا مسبقًا."""
    if os.path.exists(BIN):
        print("[build] llama-server موجود مسبقًا، تخطي البناء.")
        return

    print("[build] استنساخ llama.cpp ...")
    subprocess.run(
        ["git", "clone", "--depth", "1",
         "https://github.com/ggml-org/llama.cpp", LLAMA_DIR],
        check=True,
    )

    print("[build] تهيئة CMake ...")
    subprocess.run(
        ["cmake", "-B", "build", "-S", ".", "-DGGML_NATIVE=OFF"],
        check=True, cwd=LLAMA_DIR,
    )

    print("[build] بناء llama-server (قد يستغرق عدة دقائق) ...")
    subprocess.run(
        ["cmake", "--build", "build", "--config", "Release",
         "-j", str(os.cpu_count() or 2), "--target", "llama-server"],
        check=True, cwd=LLAMA_DIR,
    )


def download_model():
    """تنزيل ملفات GGUF من Hugging Face Hub."""
    repo = "gijl/gemma-4-E4B-it-GGUF"
    print("[download] تنزيل النموذج الأساسي ...")
    model_path = hf_hub_download(
        repo_id=repo, filename="gemma-4-E4B-it-BF16.gguf", local_dir=MODEL_DIR
    )
    print("[download] تنزيل mmproj ...")
    mmproj_path = hf_hub_download(
        repo_id=repo, filename="mmproj-BF16.gguf", local_dir=MODEL_DIR
    )
    return model_path, mmproj_path


def start_llama_server(model_path: str, mmproj_path: str):
    """تشغيل llama-server محليًا على منفذ داخلي."""
    print("[run] تشغيل llama-server ...")
    proc = subprocess.Popen([
        BIN,
        "-m", model_path,
        "--mmproj", mmproj_path,
        "--host", "127.0.0.1",
        "--port", str(LLAMA_PORT),
        "-t", "2",
        "--cache-type-k", "q8_0",
        "--cache-type-v", "iq4_nl",
        "-c", "64800",
        "-n", "50912",
    ])
    return proc


def wait_for_server(timeout_seconds: int = 180):
    print("[run] بانتظار جاهزية الخادم ...")
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        try:
            r = requests.get(f"http://127.0.0.1:{LLAMA_PORT}/health", timeout=1)
            if r.status_code == 200:
                print("[run] الخادم جاهز.")
                return
        except Exception:
            pass
        time.sleep(1)
    print("[run] تحذير: انتهت مهلة الانتظار، سيتم المتابعة رغم ذلك.")


def setup_backend():
    """هذه الدالة تقوم بتجميع وتنزيل وتشغيل النموذج، وستعمل في الخلفية."""
    global backend_ready
    try:
        build_llama_cpp()
        model_path, mmproj_path = download_model()
        start_llama_server(model_path, mmproj_path)
        wait_for_server()
        backend_ready = True
        print("[setup] اكتمل إعداد النموذج وهو جاهز الآن لاستقبال الطلبات.")
    except Exception as e:
        print(f"[setup] حدث خطأ أثناء إعداد النموذج: {e}")


# ---------------------------------------------------------------------------
# إعداد الواجهة (تطبيق FastAPI)
# ---------------------------------------------------------------------------
app = FastAPI()

@app.on_event("startup")
def startup_event():
    # تشغيل عملية التنزيل والبناء في مسار منفصل (Thread)
    # هذا يضمن أن يعمل تطبيق FastAPI فوراً دون انتظار التنزيل
    thread = threading.Thread(target=setup_backend)
    thread.start()

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"])
async def proxy(path: str, request: Request):
    # إذا لم يكتمل تنزيل النموذج بعد، يتم إرجاع رسالة توضيحية بدلاً من خطأ اتصال
    if not backend_ready:
        return JSONResponse(
            status_code=503,
            content={"error": "النموذج لا يزال قيد التنزيل والإعداد في الخلفية. يرجى المحاولة بعد قليل."}
        )

    url = f"http://127.0.0.1:{LLAMA_PORT}/{path}"
    body = await request.body()
    forward_headers = {
        k: v for k, v in request.headers.items()
        if k.lower() not in ("host", "content-length")
    }
    resp = requests.request(
        method=request.method,
        url=url,
        headers=forward_headers,
        params=request.query_params,
        data=body,
        stream=True,
    )
    response_headers = {
        k: v for k, v in resp.headers.items()
        if k.lower() != "content-encoding"
    }
    return StreamingResponse(
        resp.iter_content(chunk_size=1024),
        status_code=resp.status_code,
        headers=response_headers,
    )


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