import os

import gradio as gr
import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer


MODEL_ID = os.getenv("MODEL_ID", "trtd56/LFM2.5-1.2B-JP-bash-explainer-lora")
SYSTEM_PROMPT = "あなたは Bash コマンドの内容を説明する日本語アシスタントです。"


def load_model():
    model = AutoPeftModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype="auto",
        device_map="auto",
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer


MODEL, TOKENIZER = load_model()


def build_prompt(command: str) -> str:
    return (
        f"{SYSTEM_PROMPT}\n\n"
        "### 入力\n"
        f"{command.strip()}\n\n"
        "### 出力\n"
    )


def explain_command(command: str, max_new_tokens: int) -> str:
    command = command.strip()
    if not command:
        return "bash コマンドを入力してください。"

    prompt = build_prompt(command)
    inputs = TOKENIZER(prompt, return_tensors="pt")
    model_device = next(MODEL.parameters()).device
    inputs = {k: v.to(model_device) for k, v in inputs.items()}

    with torch.no_grad():
        output = MODEL.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=TOKENIZER.pad_token_id,
            eos_token_id=TOKENIZER.eos_token_id,
        )

    text = TOKENIZER.decode(output[0], skip_special_tokens=True)
    if "### 出力" in text:
        return text.split("### 出力", 1)[-1].strip()
    return text.strip()


EXAMPLES = [
    ["sudo apt install curl"],
    ["ls -la ~/Downloads"],
    ['grep -R "TODO" .'],
    ['find . -name "*.log" -delete'],
]

CSS = """
:root {
  --bg: #f4efe4;
  --card: #fffaf0;
  --ink: #1d1a16;
  --muted: #6b6254;
  --line: #d8ccb8;
  --accent: #9a3412;
  --accent-2: #164e63;
}
.gradio-container {
  background:
    radial-gradient(circle at top left, rgba(154, 52, 18, 0.10), transparent 28%),
    radial-gradient(circle at bottom right, rgba(22, 78, 99, 0.10), transparent 24%),
    var(--bg);
  color: var(--ink);
  font-family: "IBM Plex Sans JP", "Hiragino Sans", sans-serif;
}
.shell-card {
  border: 1px solid var(--line);
  border-radius: 20px;
  background: linear-gradient(180deg, rgba(255,250,240,0.98), rgba(255,247,235,0.98));
  box-shadow: 0 18px 60px rgba(29, 26, 22, 0.08);
}
.eyebrow {
  letter-spacing: 0.08em;
  text-transform: uppercase;
  color: var(--accent-2);
  font-size: 12px;
  font-weight: 700;
}
.hero {
  font-family: "Avenir Next", "IBM Plex Sans JP", sans-serif;
  font-size: 38px;
  line-height: 1.05;
  font-weight: 700;
  margin: 8px 0 12px;
}
.sub {
  color: var(--muted);
  font-size: 15px;
  line-height: 1.7;
}
"""


with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div class="shell-card" style="padding: 28px; margin: 24px 0 18px;">
          <div class="eyebrow">Bash Command Explainer</div>
          <div class="hero">コマンドを貼ると、日本語で説明します。</div>
          <div class="sub">
            学習済み LoRA <code>trtd56/LFM2.5-1.2B-JP-bash-explainer-lora</code> を読み込み、
            bash コマンドの動作を短い日本語で返します。
          </div>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            command = gr.Textbox(
                label="Bash コマンド",
                placeholder='例: find . -name "*.log" -delete',
                lines=5,
            )
        with gr.Column(scale=1):
            max_new_tokens = gr.Slider(
                minimum=16,
                maximum=128,
                value=64,
                step=8,
                label="最大生成トークン",
            )
            run_btn = gr.Button("説明する", variant="primary")
            clear_btn = gr.Button("クリア")

    output = gr.Textbox(label="日本語説明", lines=6, show_copy_button=True)

    gr.Examples(
        examples=EXAMPLES,
        inputs=[command],
        label="Examples",
    )

    gr.Markdown(
        "注意: ここではコマンドを実行せず、内容を説明するだけです。削除系コマンドも安全に試せます。"
    )

    run_btn.click(
        fn=explain_command,
        inputs=[command, max_new_tokens],
        outputs=output,
    )
    command.submit(
        fn=explain_command,
        inputs=[command, max_new_tokens],
        outputs=output,
    )
    clear_btn.click(lambda: ("", ""), outputs=[command, output])


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