import os
import time
import json
import shutil
import tarfile
import zipfile
import stat
import subprocess
import threading
from pathlib import Path

import requests
import gradio as gr
from huggingface_hub import hf_hub_download


APP_DIR = Path("/app")

MODEL_DIR = APP_DIR / "models"
MODEL_DIR.mkdir(parents=True, exist_ok=True)

MODEL_REPO = "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF"
MODEL_FILE = "qwen2.5-coder-3b-instruct-q4_k_m.gguf"
MODEL_PATH = MODEL_DIR / MODEL_FILE

LLAMA_DIR = APP_DIR / "llama-server"
LLAMA_EXTRACT_DIR = APP_DIR / "llama-extracted"
LLAMA_SERVER = LLAMA_DIR / "llama-server"

SERVER_HOST = "127.0.0.1"
SERVER_PORT = 8080

GRADIO_HOST = "0.0.0.0"
GRADIO_PORT = 7860

API_KEY = "local-space-key"

PARALLEL_USERS = 4
CONTEXT_PER_USER = 2048
TOTAL_CONTEXT = PARALLEL_USERS * CONTEXT_PER_USER

CPU_THREADS = int(
    os.environ.get("CPU_THREADS", "4")
)

CPU_THREADS_BATCH = int(
    os.environ.get(
        "CPU_THREADS_BATCH",
        str(CPU_THREADS)
    )
)

BATCH_SIZE = int(
    os.environ.get("BATCH_SIZE", "512")
)

MAX_TOKENS = int(
    os.environ.get("MAX_TOKENS", "512")
)

TEMPERATURE = float(
    os.environ.get("TEMPERATURE", "0.2")
)

TOP_P = float(
    os.environ.get("TOP_P", "0.9")
)

server_process = None
server_lock = threading.Lock()


def download_model():

    if MODEL_PATH.exists():

        size_gb = (
            MODEL_PATH.stat().st_size
            / 1024
            / 1024
            / 1024
        )

        print(
            f"[MODEL] Already exists: {MODEL_PATH}"
        )

        print(
            f"[MODEL] Size: {size_gb:.2f} GB"
        )

        return str(MODEL_PATH)

    print(
        "[MODEL] Downloading Qwen2.5-Coder-3B Q4_K_M..."
    )

    hf_token = os.environ.get("HF_TOKEN")

    kwargs = {
        "repo_id": MODEL_REPO,
        "filename": MODEL_FILE,
        "local_dir": str(MODEL_DIR),
    }

    if hf_token:
        kwargs["token"] = hf_token

    downloaded = hf_hub_download(
        **kwargs
    )

    print(
        f"[MODEL] Download complete: {downloaded}"
    )

    return downloaded


def find_llama_server(directory):

    directory = Path(directory)

    for path in directory.rglob(
        "llama-server"
    ):

        if path.is_file():

            try:
                path.chmod(
                    path.stat().st_mode
                    | stat.S_IXUSR
                    | stat.S_IXGRP
                    | stat.S_IXOTH
                )
            except Exception:
                pass

            return path

    return None


def find_cpu_asset():

    api_url = (
        "https://api.github.com/"
        "repos/ggml-org/llama.cpp/releases"
    )

    headers = {
        "Accept":
            "application/vnd.github+json",
        "User-Agent":
            "hf-space-qwen-llama-server"
    }

    print(
        "[LLAMA] Searching llama.cpp releases..."
    )

    response = requests.get(
        api_url,
        params={
            "per_page": 20
        },
        headers=headers,
        timeout=30
    )

    response.raise_for_status()

    releases = response.json()

    if not isinstance(
        releases,
        list
    ):
        raise RuntimeError(
            "Invalid GitHub releases response."
        )

    stable = []

    for release in releases:

        if not release.get(
            "draft",
            False
        ) and not release.get(
            "prerelease",
            False
        ):
            stable.append(release)

    releases = stable + [
        r for r in releases
        if r not in stable
    ]

    for release in releases:

        tag = release.get(
            "tag_name",
            "unknown"
        )

        assets = release.get(
            "assets",
            []
        )

        print(
            f"[LLAMA] Checking {tag}"
        )

        for asset in assets:

            name = asset.get(
                "name",
                ""
            )

            lower = name.lower()

            special = any(
                x in lower
                for x in (
                    "cuda",
                    "vulkan",
                    "rocm",
                    "openvino",
                    "sycl",
                    "metal"
                )
            )

            current_cpu = (
                "ubuntu-x64" in lower
                and lower.endswith(
                    ".tar.gz"
                )
            )

            old_cpu = (
                "ubuntu" in lower
                and "x64" in lower
                and (
                    "cpu" in lower
                    or "bin-" in lower
                )
                and lower.endswith(
                    (
                        ".tar.gz",
                        ".zip"
                    )
                )
            )

            if (
                not special
                and (
                    current_cpu
                    or old_cpu
                )
            ):

                print(
                    f"[LLAMA] Found CPU build: {name}"
                )

                return asset

    print(
        "[LLAMA] No compatible CPU build found."
    )

    for release in releases[:5]:

        print(
            "Release:",
            release.get(
                "tag_name",
                "unknown"
            )
        )

        for asset in release.get(
            "assets",
            []
        ):

            print(
                " ",
                asset.get(
                    "name",
                    ""
                )
            )

    raise RuntimeError(
        "Could not find Ubuntu x64 CPU "
        "llama.cpp release asset."
    )


def download_llama_server():

    if LLAMA_SERVER.exists():

        try:
            LLAMA_SERVER.chmod(
                LLAMA_SERVER.stat().st_mode
                | stat.S_IXUSR
                | stat.S_IXGRP
                | stat.S_IXOTH
            )
        except Exception:
            pass

        print(
            f"[LLAMA] Existing server found: "
            f"{LLAMA_SERVER}"
        )

        return str(
            LLAMA_SERVER
        )

    print(
        "[LLAMA] Downloading prebuilt llama.cpp server..."
    )

    asset = find_cpu_asset()

    asset_name = asset[
        "name"
    ]

    download_url = asset[
        "browser_download_url"
    ]

    archive_path = (
        APP_DIR /
        asset_name
    )

    print(
        f"[LLAMA] Selected: {asset_name}"
    )

    if not archive_path.exists():

        print(
            "[LLAMA] Downloading archive..."
        )

        headers = {
            "User-Agent":
                "hf-space-qwen-llama-server"
        }

        with requests.get(
            download_url,
            headers=headers,
            stream=True,
            timeout=120
        ) as response:

            response.raise_for_status()

            total = int(
                response.headers.get(
                    "content-length",
                    0
                )
            )

            downloaded = 0
            last_print = 0

            with open(
                archive_path,
                "wb"
            ) as f:

                for chunk in response.iter_content(
                    chunk_size=1024 * 1024
                ):

                    if not chunk:
                        continue

                    f.write(chunk)

                    downloaded += len(
                        chunk
                    )

                    if total:

                        now = time.time()

                        if now - last_print >= 0.5:

                            percent = (
                                downloaded
                                * 100
                                / total
                            )

                            print(
                                f"\r[LLAMA] "
                                f"{percent:5.1f}% "
                                f"{downloaded / 1024 / 1024:.1f} MB / "
                                f"{total / 1024 / 1024:.1f} MB",
                                end="",
                                flush=True
                            )

                            last_print = now

        print()

        print(
            "[LLAMA] Archive download complete."
        )

    else:

        print(
            "[LLAMA] Archive already exists."
        )

    if LLAMA_EXTRACT_DIR.exists():

        shutil.rmtree(
            LLAMA_EXTRACT_DIR
        )

    LLAMA_EXTRACT_DIR.mkdir(
        parents=True,
        exist_ok=True
    )

    print(
        "[LLAMA] Extracting..."
    )

    if asset_name.lower().endswith(
        ".tar.gz"
    ):

        with tarfile.open(
            archive_path,
            "r:gz"
        ) as tar:

            tar.extractall(
                LLAMA_EXTRACT_DIR
            )

    elif asset_name.lower().endswith(
        ".zip"
    ):

        with zipfile.ZipFile(
            archive_path,
            "r"
        ) as z:

            z.extractall(
                LLAMA_EXTRACT_DIR
            )

    else:

        raise RuntimeError(
            f"Unsupported archive: {asset_name}"
        )

    found = find_llama_server(
        LLAMA_EXTRACT_DIR
    )

    if found is None:

        for path in LLAMA_EXTRACT_DIR.rglob(
            "*"
        ):

            if path.is_file():
                print(
                    path
                )

        raise RuntimeError(
            "llama-server was not found "
            "inside the archive."
        )

    print(
        f"[LLAMA] Found: {found}"
    )

    source_dir = found.parent

    LLAMA_DIR.mkdir(
        parents=True,
        exist_ok=True
    )

    for item in source_dir.iterdir():

        destination = (
            LLAMA_DIR /
            item.name
        )

        if item.is_dir():

            if destination.exists():
                shutil.rmtree(
                    destination
                )

            shutil.copytree(
                item,
                destination
            )

        else:

            shutil.copy2(
                item,
                destination
            )

    final_server = (
        LLAMA_DIR /
        found.name
    )

    if not final_server.exists():

        raise RuntimeError(
            "Failed to install llama-server."
        )

    final_server.chmod(
        final_server.stat().st_mode
        | stat.S_IXUSR
        | stat.S_IXGRP
        | stat.S_IXOTH
    )

    library_dirs = []

    for root, dirs, files in os.walk(
        LLAMA_DIR
    ):

        for filename in files:

            if (
                filename.endswith(".so")
                or ".so." in filename
            ):

                library_dirs.append(
                    root
                )

                break

    library_dirs = list(
        dict.fromkeys(
            library_dirs
        )
    )

    if library_dirs:

        old_ld = os.environ.get(
            "LD_LIBRARY_PATH",
            ""
        )

        new_ld = ":".join(
            library_dirs
        )

        if old_ld:

            os.environ[
                "LD_LIBRARY_PATH"
            ] = (
                new_ld
                + ":"
                + old_ld
            )

        else:

            os.environ[
                "LD_LIBRARY_PATH"
            ] = new_ld

        print(
            "[LLAMA] Libraries configured."
        )

    print(
        f"[LLAMA] Server installed: "
        f"{final_server}"
    )

    return str(
        final_server
    )


def start_llama_server():

    global server_process

    with server_lock:

        if server_process is not None:

            if (
                server_process.poll()
                is None
            ):

                return

        server_path = (
            download_llama_server()
        )

        command = [
            server_path,
            "-m",
            str(MODEL_PATH),
            "-np",
            str(PARALLEL_USERS),
            "-c",
            str(TOTAL_CONTEXT),
            "--cont-batching",
            "-ngl",
            "0",
            "-t",
            str(CPU_THREADS),
            "-tb",
            str(CPU_THREADS_BATCH),
            "-b",
            str(BATCH_SIZE),
            "--host",
            SERVER_HOST,
            "--port",
            str(SERVER_PORT),
            "--api-key",
            API_KEY
        ]

        print(
            "[LLAMA] Starting server..."
        )

        server_process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            env=os.environ.copy()
        )

        def read_logs():

            try:

                for line in iter(
                    server_process.stdout.readline,
                    ""
                ):

                    if line:
                        print(
                            "[llama.cpp]",
                            line.rstrip()
                        )

            except Exception:
                pass

        threading.Thread(
            target=read_logs,
            daemon=True
        ).start()

    health_url = (
        f"http://{SERVER_HOST}:"
        f"{SERVER_PORT}/health"
    )

    print(
        "[LLAMA] Waiting for server..."
    )

    started = time.time()

    while True:

        if (
            server_process is not None
            and server_process.poll()
            is not None
        ):

            raise RuntimeError(
                "llama-server stopped "
                f"with code "
                f"{server_process.returncode}"
            )

        try:

            response = requests.get(
                health_url,
                timeout=2
            )

            if response.status_code in (
                200,
                204
            ):

                print(
                    "[LLAMA] Server is ready."
                )

                return

        except Exception:
            pass

        if (
            time.time() - started
            > 180
        ):

            raise RuntimeError(
                "Timed out waiting for llama-server."
            )

        time.sleep(1)


def clean_content(content):

    if content is None:
        return ""

    if isinstance(
        content,
        str
    ):
        return content

    if isinstance(
        content,
        list
    ):

        result = []

        for item in content:

            if isinstance(
                item,
                str
            ):

                result.append(
                    item
                )

            elif isinstance(
                item,
                dict
            ):

                if "text" in item:

                    result.append(
                        str(
                            item["text"]
                        )
                    )

                elif "content" in item:

                    result.append(
                        str(
                            item["content"]
                        )
                    )

            else:

                result.append(
                    str(item)
                )

        return "\n".join(
            x for x in result
            if x
        )

    if isinstance(
        content,
        dict
    ):

        if "text" in content:

            return str(
                content["text"]
            )

        if "content" in content:

            return str(
                content["content"]
            )

        return json.dumps(
            content,
            ensure_ascii=False
        )

    return str(content)


def build_messages(
    message,
    history
):

    messages = []

    history = history or []

    recent_history = history[-10:]

    for item in recent_history:

        if not isinstance(
            item,
            dict
        ):
            continue

        role = item.get(
            "role"
        )

        content = clean_content(
            item.get(
                "content",
                ""
            )
        )

        if role in (
            "user",
            "assistant",
            "system"
        ) and content:

            messages.append({
                "role":
                    role,
                "content":
                    content
            })

    messages.append({
        "role":
            "user",
        "content":
            clean_content(
                message
            )
    })

    return messages


def generate_response(
    message,
    history
):

    if not message:
        return "Please enter a message."

    messages = build_messages(
        message,
        history
    )

    url = (
        f"http://{SERVER_HOST}:"
        f"{SERVER_PORT}/v1/chat/completions"
    )

    payload = {
        "model":
            "qwen2.5-coder-3b",
        "messages":
            messages,
        "temperature":
            TEMPERATURE,
        "top_p":
            TOP_P,
        "max_tokens":
            MAX_TOKENS,
        "stream":
            False
    }

    headers = {
        "Content-Type":
            "application/json",
        "Authorization":
            f"Bearer {API_KEY}"
    }

    try:

        response = requests.post(
            url,
            headers=headers,
            json=payload,
            timeout=300
        )

    except requests.RequestException as e:

        return (
            "Server connection error:\n\n"
            f"{e}"
        )

    if response.status_code != 200:

        try:

            data = response.json()

            return (
                "llama.cpp error:\n\n"
                + json.dumps(
                    data,
                    indent=2,
                    ensure_ascii=False
                )
            )

        except Exception:

            return (
                f"llama.cpp HTTP "
                f"{response.status_code}\n\n"
                f"{response.text}"
            )

    try:

        data = response.json()

        choices = data.get(
            "choices",
            []
        )

        if not choices:

            return "No response generated."

        choice = choices[0]

        message_data = choice.get(
            "message",
            {}
        )

        result = message_data.get(
            "content",
            ""
        )

        if not result:

            result = choice.get(
                "text",
                ""
            )

        return result.strip()

    except Exception as e:

        return (
            "Response parsing error:\n\n"
            f"{e}\n\n"
            f"{response.text}"
        )


def create_ui():

    demo = gr.ChatInterface(
        fn=generate_response,
        concurrency_limit=PARALLEL_USERS
    )

    return demo


def initialize():

    print(
        "============================================================"
    )

    print(
        "Qwen2.5-Coder-3B llama.cpp server"
    )

    print(
        "============================================================"
    )

    print(
        f"[CONFIG] Users: {PARALLEL_USERS}"
    )

    print(
        f"[CONFIG] Context/user: {CONTEXT_PER_USER}"
    )

    print(
        f"[CONFIG] CPU threads: {CPU_THREADS}"
    )

    print(
        f"[CONFIG] Batch threads: {CPU_THREADS_BATCH}"
    )

    print(
        f"[CONFIG] Batch size: {BATCH_SIZE}"
    )

    print(
        f"[CONFIG] Max tokens: {MAX_TOKENS}"
    )

    print(
        "============================================================"
    )

    download_model()

    download_llama_server()

    start_llama_server()

    print(
        "============================================================"
    )

    print(
        "INITIALIZATION COMPLETE"
    )

    print(
        f"Model: {MODEL_FILE}"
    )

    print(
        f"Users: {PARALLEL_USERS}"
    )

    print(
        f"Context/user: {CONTEXT_PER_USER}"
    )

    print(
        f"CPU threads: {CPU_THREADS}"
    )

    print(
        "============================================================"
    )


def shutdown():

    global server_process

    print(
        "[SYSTEM] Shutting down..."
    )

    if server_process is None:
        return

    try:

        if (
            server_process.poll()
            is None
        ):

            server_process.terminate()

            try:

                server_process.wait(
                    timeout=5
                )

            except subprocess.TimeoutExpired:

                server_process.kill()

    except Exception as e:

        print(
            f"[SYSTEM] Shutdown error: {e}"
        )


if __name__ == "__main__":

    try:

        initialize()

        demo = create_ui()

        demo.queue(
            default_concurrency_limit=
                PARALLEL_USERS
        )

        demo.launch(
            server_name=
                GRADIO_HOST,
            server_port=
                GRADIO_PORT,
            show_error=True,
            max_threads=
                PARALLEL_USERS + 4,
            share=False
        )

    except KeyboardInterrupt:

        print(
            "[SYSTEM] Interrupted."
        )

    except Exception as e:

        print(
            "============================================================"
        )

        print(
            "STARTUP FAILED"
        )

        print(
            repr(e)
        )

        print(
            "============================================================"
        )

        raise

    finally:

        shutdown()