import time

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama


# -----------------------------
# 1. Measure model startup time
# -----------------------------

startup_start = time.perf_counter()

print("Downloading/locating model from Hugging Face...")

model_path = hf_hub_download(
    repo_id="iamabhayaditya/EfficientMath-AI",
    filename="Meta-Llama-3.1-8B.Q4_K_M.gguf",
)

download_finished = time.perf_counter()

print("Loading model into memory...")

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=4,
)

model_loaded_finished = time.perf_counter()

download_time = download_finished - startup_start
model_load_time = model_loaded_finished - download_finished
total_startup_time = model_loaded_finished - startup_start

print(f"Model download time: {download_time:.2f} seconds")
print(f"Model loading time: {model_load_time:.2f} seconds")
print(f"Total startup time: {total_startup_time:.2f} seconds")
print("Model loaded successfully.")


# Used to identify the first request after the model is loaded
request_count = 0


# -----------------------------
# 2. Prediction function
# -----------------------------

def solve_math_problem(question):
    global request_count

    request_count += 1
    request_number = request_count

    prompt = (
        "Below is a math word problem. Solve it step by step "
        "and provide the final answer.\n\n"
        f"### Problem:\n{question}\n\n"
        "### Solution:\n"
    )

    request_start = time.perf_counter()
    first_output_time = None
    generated_text = ""

    try:
        stream = llm(
            prompt,
            max_tokens=256,
            temperature=0.2,
            top_p=0.9,
            stream=True,
            stop=[
                "<|end_of_text|>",
                "</s>",
                "<|eot_id|>",
            ],
        )

        for output in stream:
            new_text = output["choices"][0]["text"]

            if first_output_time is None:
                first_output_time = time.perf_counter()

            generated_text += new_text
            yield generated_text

        request_finished = time.perf_counter()

        total_latency = request_finished - request_start

        if first_output_time is not None:
            time_to_first_output = first_output_time - request_start
        else:
            time_to_first_output = None

        # Count generated tokens after generation is complete
        generated_tokens = len(
            llm.tokenize(
                generated_text.encode("utf-8"),
                add_bos=False,
            )
        )

        if total_latency > 0:
            tokens_per_second = generated_tokens / total_latency
        else:
            tokens_per_second = 0

        request_type = (
            "FIRST REQUEST AFTER MODEL LOAD"
            if request_number == 1
            else "WARM REQUEST"
        )

        print("\n----- INFERENCE METRICS -----")
        print(f"Request type: {request_type}")
        print(f"Request number: {request_number}")
        print(f"Total latency: {total_latency:.2f} seconds")
        print(
            f"Time to first output: "
            f"{time_to_first_output:.2f} seconds"
            if time_to_first_output is not None
            else "Time to first output: unavailable"
        )
        print(f"Generated tokens: {generated_tokens}")
        print(f"Throughput: {tokens_per_second:.2f} tokens/second")
        print("-----------------------------\n")

    except Exception as error:
        print(f"Server error: {error}")

        yield (
            "I encountered an issue while solving the problem. "
            "Please try again."
        )


# -----------------------------
# 3. Gradio interface
# -----------------------------

custom_css = """
.gradio-container {
    background-color: #000000 !important;
}

.markdown-text h1 {
    color: #ff7f00 !important;
}

.markdown-text p {
    color: #cccccc !important;
}

textarea {
    border: 2px solid #ff7f00 !important;
    background-color: #111111 !important;
    color: #ffffff !important;
}

button.primary {
    background: linear-gradient(90deg, #ff7f00, #ffaa00) !important;
    border: none !important;
    color: black !important;
    font-weight: bold !important;
}

span.svelte-1gfkn6j,
.label {
    color: #ff7f00 !important;
}
"""


with gr.Blocks(
    theme=gr.themes.Monochrome(),
    css=custom_css,
) as app:
    gr.Markdown(
        "<h1 style='text-align: center; margin-top: 20px;'>"
        "EfficientMath-AI"
        "</h1>"
    )

    gr.Markdown(
        "<p style='text-align: center;'>"
        "This is a custom fine-tuned Llama 3.1 8B model, "
        "trained to solve grade school math word problems."
        "</p>"
    )

    with gr.Row():
        with gr.Column(scale=1):
            user_input = gr.Textbox(
                lines=5,
                placeholder="Enter a math word problem here...",
                label="Question",
            )

            gr.Examples(
                examples=[
                    (
                        "A bag containing 30 apples weighs 6 kg. "
                        "How much will 1080 apples weigh?"
                    ),
                    (
                        "If the cost of 18 apples is 90 rupees, "
                        "what is the cost of 24 apples?"
                    ),
                    (
                        "Abhay has 16 apples, he borrowed 5 from Akash "
                        "then gave 14 to Shivam. "
                        "How many apples is Abhay left with?"
                    ),
                ],
                inputs=user_input,
                label="Click an example below to test:",
            )

            with gr.Row():
                clear_btn = gr.ClearButton([user_input])
                submit_btn = gr.Button(
                    "Submit",
                    variant="primary",
                )

        with gr.Column(scale=1):
            model_output = gr.Textbox(
                label="Model Solution",
                lines=5,
                max_lines=50,
                interactive=False,
            )

    submit_btn.click(
        fn=solve_math_problem,
        inputs=user_input,
        outputs=model_output,
    )


# -----------------------------
# 4. Launch Hugging Face Space
# -----------------------------

app.launch()