import os
import subprocess
import signal
import tempfile
from pathlib import Path
from textwrap import dedent
import logging
import gradio as gr
from huggingface_hub import HfApi, ModelCard, whoami
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from apscheduler.schedulers.background import BackgroundScheduler
from datetime import datetime
import numpy as np

os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
CONVERSION_SCRIPT = "./llama.cpp/convert_hf_to_gguf.py"
logger = logging.getLogger(__name__)

def get_repo_namespace(repo_owner, username, user_orgs):
    if repo_owner == 'self':
        return username
    for org in user_orgs:
        if org['name'] == repo_owner:
            return org['name']
    raise ValueError(f"Invalid repo_owner: {repo_owner}")

def escape(s: str) -> str:
    return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("\n", "<br/>")

def toggle_repo_owner(export_to_org, oauth_token: gr.OAuthToken | None):
    if oauth_token is None or oauth_token.token is None:
        raise gr.Error("You must be logged in to use GGUF-my-repo")
    if not export_to_org:
        return gr.update(visible=False, choices=["self"], value="self"), gr.update(visible=False, value="")
    info = whoami(oauth_token.token)
    orgs = [org["name"] for org in info.get("orgs", [])]
    return gr.update(visible=True, choices=["self"] + orgs, value="self"), gr.update(visible=True)
def generate_importance_matrix(model_path: str, train_data_path: str, output_path: str):
    imatrix_command = [
        "./llama.cpp/llama-imatrix",
        "-m", model_path,
        "-f", train_data_path,
        "-ngl", "99",
        "--output-frequency", "10",
        "-o", output_path,
    ]

    if not os.path.isfile(model_path):
        raise Exception(f"Model file not found: {model_path}")

    print("Running imatrix command...")
    process = subprocess.Popen(imatrix_command, shell=False)

    try:
        process.wait(timeout=60)  # added wait
    except subprocess.TimeoutExpired:
        print("Imatrix computation timed out. Sending SIGINT to allow graceful termination...")
        process.send_signal(signal.SIGINT)
        try:
            process.wait(timeout=5)  # grace period
        except subprocess.TimeoutExpired:
            print("Imatrix proc still didn't term. Forecfully terming process...")
            process.kill()

    print("Importance matrix generation completed.")

def split_upload_model(model_path: str, outdir: str, repo_id: str, oauth_token: gr.OAuthToken | None, split_max_tensors=256, split_max_size=None, org_token=None, export_to_org=False):
    print(f"Model path: {model_path}")
    print(f"Output dir: {outdir}")

    if oauth_token is None or oauth_token.token is None:
        raise ValueError("You have to be logged in.")

    split_cmd = [
        "./llama.cpp/llama-gguf-split",
        "--split",
    ]
    if split_max_size:
        split_cmd.append("--split-max-size")
        split_cmd.append(split_max_size)
    else:
        split_cmd.append("--split-max-tensors")
        split_cmd.append(str(split_max_tensors))

    # args for output
    model_path_prefix = '.'.join(model_path.split('.')[:-1]) # remove the file extension
    split_cmd.append(model_path)
    split_cmd.append(model_path_prefix)

    print(f"Split command: {split_cmd}")

    result = subprocess.run(split_cmd, shell=False, capture_output=True, text=True)
    print(f"Split command stdout: {result.stdout}")
    print(f"Split command stderr: {result.stderr}")

    if result.returncode != 0:
        stderr_str = result.stderr.decode("utf-8")
        raise Exception(f"Error splitting the model: {stderr_str}")
    print("Model split successfully!")

    # remove the original model file if needed
    if os.path.exists(model_path):
        os.remove(model_path)

    model_file_prefix = model_path_prefix.split('/')[-1]
    print(f"Model file name prefix: {model_file_prefix}")
    sharded_model_files = [f for f in os.listdir(outdir) if f.startswith(model_file_prefix) and f.endswith(".gguf")]
    if sharded_model_files:
        print(f"Sharded model files: {sharded_model_files}")
        if export_to_org and org_token!="":
          api = HfApi(token = org_token)
        else:
          api = HfApi(token=oauth_token.token)
        for file in sharded_model_files:
            file_path = os.path.join(outdir, file)
            print(f"Uploading file: {file_path}")
            try:
                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=file,
                    repo_id=repo_id,
                )
            except Exception as e:
                raise Exception(f"Error uploading file {file_path}: {e}")
    else:
        raise Exception("No sharded files found.")

    print("Sharded model has been uploaded successfully!")

def process_model(model_id, q_method, use_imatrix, imatrix_q_method, private_repo,
                  train_data_file, split_model, split_max_tensors, split_max_size,
                  export_to_org, repo_owner, org_token, oauth_token: gr.OAuthToken | None):
    if oauth_token is None or oauth_token.token is None:
        raise gr.Error("You must be logged in to use GGUF-my-repo")

    user_info = whoami(oauth_token.token)
    username = user_info["name"]
    user_orgs = user_info.get("orgs", [])
    if not export_to_org:
        repo_owner = "self"


    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"Time {current_time}, Username {username}, Model_ID, {model_id}, q_method {','.join(q_method)}")

    repo_namespace = get_repo_namespace(repo_owner, username, user_orgs)
    model_name = model_id.split('/')[-1]
    api_token = org_token if (export_to_org and org_token!="") else oauth_token.token
    api = HfApi(token=api_token)

    dl_pattern = ["*.md", "*.json", "*.model"]
    pattern = "*.safetensors" if any(
        f.path.endswith(".safetensors")
        for f in api.list_repo_tree(repo_id=model_id, recursive=True)
    ) else "*.bin"
    dl_pattern += [pattern]

    os.makedirs("downloads", exist_ok=True)
    os.makedirs("outputs", exist_ok=True)

    with tempfile.TemporaryDirectory(dir="outputs") as outdir:
        fp16 = str(Path(outdir)/f"{model_name}.fp16.gguf")

        with tempfile.TemporaryDirectory(dir="downloads") as tmpdir:
            local_dir = Path(tmpdir)/model_name
            api.snapshot_download(repo_id=model_id, local_dir=local_dir, local_dir_use_symlinks=False, allow_patterns=dl_pattern)

            config_dir = local_dir/"config.json"
            adapter_config_dir = local_dir/"adapter_config.json"
            if os.path.exists(adapter_config_dir) and not os.path.exists(config_dir):
                raise Exception("adapter_config.json is present. If converting LoRA, use GGUF-my-lora.")

            result = subprocess.run(["python", CONVERSION_SCRIPT, local_dir, "--outtype", "f16", "--outfile", fp16], shell=False, capture_output=True)
            if result.returncode != 0:
                raise Exception(f"Error converting to fp16: {result.stderr.decode()}")

        imatrix_path = Path(outdir)/"imatrix.dat"
        if use_imatrix:
            train_data_path = train_data_file.name if train_data_file else "llama.cpp/groups_merged.txt"
            if not os.path.isfile(train_data_path):
                raise Exception(f"Training data not found: {train_data_path}")
            generate_importance_matrix(fp16, train_data_path, imatrix_path)

        quant_methods = [imatrix_q_method] if use_imatrix else (q_method if isinstance(q_method, list) else [q_method])
        suffix = "imat" if use_imatrix else None

        gguf_files = []
        for method in quant_methods:
            name = f"{model_name.lower()}-{method.lower()}-{suffix}.gguf" if suffix else f"{model_name.lower()}-{method.lower()}.gguf"
            path = str(Path(outdir)/name)
            quant_cmd = ["./llama.cpp/llama-quantize", "--imatrix", imatrix_path, fp16, path, method] if use_imatrix else ["./llama.cpp/llama-quantize", fp16, path, method]
            result = subprocess.run(quant_cmd, shell=False, capture_output=True)
            if result.returncode != 0:
                raise Exception(f"Quantization failed ({method}): {result.stderr.decode()}")
            gguf_files.append((name, path))

        suffix_for_repo = f"{imatrix_q_method}-imat" if use_imatrix else "-".join(quant_methods)
        repo_id = f"{repo_namespace}/{model_name}-{suffix_for_repo}-GGUF"
        new_repo_url = api.create_repo(repo_id=repo_id, exist_ok=True, private=private_repo)

        try:
            card = ModelCard.load(model_id, token=oauth_token.token)
        except:
            card = ModelCard("")
        card.data.tags = (card.data.tags or []) + ["llama-cpp", "gguf-my-repo"]
        card.data.base_model = model_id
        card.text = dedent(f"""
            # {repo_id}
            Absolutely tremendous! This repo features **GGUF quantized** versions of [{model_id}](https://huggingface.co/{model_id}) — made possible using the *very powerful* `llama.cpp`. Believe me, it's fast, it's smart, it's winning.
            ## Quantized Versions:
            Only the best quantization. You’ll love it.
            ## Run with llama.cpp
            Just plug it in, hit the command line, and boom — you're running world-class AI, folks:
            ```bash
            llama-cli --hf-repo {repo_id} --hf-file {gguf_files[0][0]} -p "AI First, but also..."
            ```
            This beautiful Hugging Face Space was brought to you by the **amazing team at [Antigma Labs](https://antigma.ai)**. Great people. Big vision. Doing things that matter — and doing them right.
            Total winners.
        """)
        readme_path = Path(outdir)/"README.md"
        card.save(readme_path)
        for name, path in gguf_files:
            if split_model:
                split_upload_model(path, outdir, repo_id, oauth_token, split_max_tensors, split_max_size, org_token, export_to_org)
            else:
                api.upload_file(path_or_fileobj=path, path_in_repo=name, repo_id=repo_id)
        if use_imatrix and os.path.isfile(imatrix_path):
            api.upload_file(path_or_fileobj=imatrix_path, path_in_repo="imatrix.dat", repo_id=repo_id)
        api.upload_file(path_or_fileobj=readme_path, path_in_repo="README.md", repo_id=repo_id)

        return (f'<h1>✅ DONE</h1><br/>Repo: <a href="{new_repo_url}" target="_blank" style="text-decoration:underline">{repo_id}</a>', f"llama{np.random.randint(9)}.png")


css="""/* Custom CSS to allow scrolling */
.gradio-container {overflow-y: auto;}
"""
model_id = HuggingfaceHubSearch(
    label="Hub Model ID",
    placeholder="Search for model id on Huggingface",
    search_type="model",
)

export_to_org = gr.Checkbox(
    label="Export to Organization Repository",
    value=False,
    info="If checked, you can select an organization to export to."
)

repo_owner = gr.Dropdown(
    choices=["self"],
    value="self",
    label="Repository Owner",
    visible=False
)

org_token = gr.Textbox(
    label="Org Access Token",
    type="password",
    visible=False
)

q_method = gr.Dropdown(
    ["Q2_K", "Q3_K_S", "Q3_K_M", "Q3_K_L", "Q4_0", "Q4_K_S", "Q4_K_M", "Q5_0", "Q5_K_S", "Q5_K_M", "Q6_K", "Q8_0"],
    label="Quantization Method",
    info="GGML quantization type",
    value="Q4_K_M",
    filterable=False,
    visible=True,
    multiselect=True
)

imatrix_q_method = gr.Dropdown(
    ["IQ3_M", "IQ3_XXS", "Q4_K_M", "Q4_K_S", "IQ4_NL", "IQ4_XS", "Q5_K_M", "Q5_K_S"],
    label="Imatrix Quantization Method",
    info="GGML imatrix quants type",
    value="IQ4_NL",
    filterable=False,
    visible=False
)

use_imatrix = gr.Checkbox(
    value=False,
    label="Use Imatrix Quantization",
    info="Use importance matrix for quantization."
)

private_repo = gr.Checkbox(
    value=False,
    label="Private Repo",
    info="Create a private repo under your username."
)

train_data_file = gr.File(
    label="Training Data File",
    file_types=["txt"],
    visible=False
)

split_model = gr.Checkbox(
    value=False,
    label="Split Model",
    info="Shard the model using gguf-split."
)

split_max_tensors = gr.Number(
    value=256,
    label="Max Tensors per File",
    info="Maximum number of tensors per file when splitting model.",
    visible=False
)

split_max_size = gr.Textbox(
    label="Max File Size",
    info="Maximum file size when splitting model (--split-max-size). May leave empty to use the default. Accepted suffixes: M, G. Example: 256M, 5G",
    visible=False
)

iface = gr.Interface(
    fn=process_model,
    inputs=[
        model_id,
        q_method,
        use_imatrix,
        imatrix_q_method,
        private_repo,
        train_data_file,
        split_model,
        split_max_tensors,
        split_max_size,
        export_to_org,
        repo_owner,
        org_token
    ],
    outputs=[
        gr.Markdown(label="Output"),
        gr.Image(show_label=False)
    ],
    title="Make your own GGUF Quants — faster than ever before, believe me.",
    description="We take your Hugging Face repo — a terrific repo — we quantize it, we package it beautifully, and we give you your very own repo. It's smart. It's efficient. It's huge. You're gonna love it.",
    api_name=False
)
with gr.Blocks(css=".gradio-container {overflow-y: auto;}",theme=gr.themes.Glass()) as demo:
    gr.Markdown("Logged in, you must be. Classy, secure, and victorious, it keeps us.")
    gr.LoginButton(min_width=250)



    export_to_org.change(fn=toggle_repo_owner, inputs=[export_to_org], outputs=[repo_owner, org_token])

    split_model.change(fn=lambda sm: (gr.update(visible=sm), gr.update(visible=sm)), inputs=split_model, outputs=[split_max_tensors, split_max_size])
    use_imatrix.change(fn=lambda use: (gr.update(visible=not use), gr.update(visible=use), gr.update(visible=use)), inputs=use_imatrix, outputs=[q_method, imatrix_q_method, train_data_file])

    iface.render()



def restart_space():
    HfApi().restart_space(repo_id="Antigma/quantize-my-repo", token=HF_TOKEN, factory_reboot=True)

scheduler = BackgroundScheduler()
scheduler.add_job(restart_space, "interval", seconds=21600)
scheduler.start()

demo.queue(default_concurrency_limit=1, max_size=5).launch(debug=True, show_api=False)