"""Hugging Face Space app for HALT-CoT."""

from __future__ import annotations

import os

import gradio as gr

from halt_cot import HaltCoTConfig
from halt_cot.transformers_backend import HaltCoTForCausalLM


DEFAULT_MODEL_ID = os.getenv("HALT_COT_MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")
DEFAULT_DEVICE_MAP = os.getenv("HALT_COT_DEVICE_MAP") or None
DEFAULT_THETA = float(os.getenv("HALT_COT_THETA", "0.6"))

_RUNNER: HaltCoTForCausalLM | None = None


def get_runner() -> HaltCoTForCausalLM:
    global _RUNNER
    if _RUNNER is None:
        _RUNNER = HaltCoTForCausalLM.from_pretrained(
            DEFAULT_MODEL_ID,
            device_map=DEFAULT_DEVICE_MAP,
        )
    return _RUNNER


def parse_candidates(candidate_text: str) -> list[str]:
    lines = []
    for raw_line in candidate_text.replace(",", "\n").splitlines():
        line = raw_line.strip()
        if line:
            lines.append(line)
    if len(lines) < 2:
        raise gr.Error("Provide at least two candidate answers.")
    return lines


def run_halt_cot(
    question: str,
    candidate_text: str,
    theta: float,
    consecutive: int,
    max_steps: int,
    step_max_new_tokens: int,
):
    if not question.strip():
        raise gr.Error("Question is required.")

    config = HaltCoTConfig(
        theta=float(theta),
        consecutive_low_entropy=int(consecutive),
        max_steps=int(max_steps),
        step_max_new_tokens=int(step_max_new_tokens),
    )
    result = get_runner().run(question.strip(), parse_candidates(candidate_text), config=config)
    trace_rows = [
        [
            step.index,
            round(step.entropy, 4),
            step.prediction,
            step.generated_tokens,
            step.halted,
            step.text,
        ]
        for step in result.steps
    ]
    final_probs = []
    if result.steps:
        final_probs = sorted(
            result.steps[-1].probabilities.items(),
            key=lambda item: item[1],
            reverse=True,
        )
        final_probs = [[label, round(probability, 6)] for label, probability in final_probs]

    status = "halted" if result.halted else "reached max steps"
    summary = (
        f"Answer: {result.answer}\n"
        f"Status: {status}\n"
        f"Generated reasoning tokens: {result.generated_tokens}"
    )
    return summary, result.reasoning, trace_rows, final_probs


with gr.Blocks(title="HALT-CoT") as demo:
    gr.Markdown("# HALT-CoT")
    with gr.Row():
        with gr.Column(scale=2):
            question = gr.Textbox(
                label="Question",
                lines=5,
                value="If a shop has 12 apples and sells 5, how many apples are left?",
            )
            candidates = gr.Textbox(
                label="Candidate answers",
                lines=6,
                value="5\n6\n7\n8\n9",
            )
            run_button = gr.Button("Run HALT-CoT", variant="primary")
        with gr.Column(scale=1):
            theta = gr.Slider(0.0, 2.0, value=DEFAULT_THETA, step=0.05, label="Entropy threshold")
            consecutive = gr.Slider(1, 4, value=2, step=1, label="Consecutive low-entropy steps")
            max_steps = gr.Slider(1, 12, value=8, step=1, label="Maximum steps")
            step_max_new_tokens = gr.Slider(8, 160, value=64, step=8, label="Step token cap")

    summary = gr.Textbox(label="Result", lines=3)
    reasoning = gr.Textbox(label="Generated reasoning", lines=8)
    trace = gr.Dataframe(
        headers=["Step", "Entropy", "Prediction", "Tokens", "Halted", "Text"],
        label="HALT trace",
    )
    probabilities = gr.Dataframe(headers=["Candidate", "Probability"], label="Final answer distribution")

    run_button.click(
        run_halt_cot,
        inputs=[question, candidates, theta, consecutive, max_steps, step_max_new_tokens],
        outputs=[summary, reasoning, trace, probabilities],
    )


if __name__ == "__main__":
    demo.launch()