```python
# app.py
import time
import zipfile
from pathlib import Path

import gradio as gr

from src.train import finetune_lora
from src.infer import load_generator, generate_text


def _default_output_root() -> Path:
    # On Hugging Face Spaces, /data exists if Persistent Storage is enabled.
    # Otherwise fall back to a repo-local outputs/ directory.
    return Path("/data/outputs") if Path("/data").exists() else Path("outputs")


def _zip_dir(src_dir: Path, zip_path: Path) -> Path:
    zip_path.parent.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for p in src_dir.rglob("*"):
            if p.is_file():
                zf.write(p, arcname=p.relative_to(src_dir))
    return zip_path


def run_train(
    base_model: str,
    dataset_id: str,
    max_train_samples: int,
    max_steps: int,
    lr: float,
    batch_size: int,
    lora_r: int,
    lora_alpha: int,
    lora_dropout: float,
):
    out_root = _default_output_root()
    run_id = time.strftime("%Y%m%d-%H%M%S")
    out_dir = out_root / run_id
    out_dir.mkdir(parents=True, exist_ok=True)

    status = finetune_lora(
        base_model=base_model.strip(),
        dataset_id=dataset_id.strip(),
        output_dir=str(out_dir),
        max_train_samples=int(max_train_samples),
        max_steps=int(max_steps),
        learning_rate=float(lr),
        batch_size=int(batch_size),
        lora_r=int(lora_r),
        lora_alpha=int(lora_alpha),
        lora_dropout=float(lora_dropout),
    )

    adapter_dir = out_dir / "adapter"
    zip_path = out_dir / "adapter.zip"

    zip_file = None
    if adapter_dir.exists():
        _zip_dir(adapter_dir, zip_path)
        zip_file = str(zip_path)

    msg = (
        f"Done.\n\n"
        f"Run dir: {out_dir}\n"
        f"Adapter dir: {adapter_dir}\n\n"
        f"{status}"
    )

    return msg, zip_file, str(out_dir), str(adapter_dir)


def run_generate(
    base_model: str,
    adapter_dir: str,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
):
    gen = load_generator(base_model.strip(), adapter_dir.strip())
    return generate_text(
        gen,
        prompt,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
    )


with gr.Blocks(title="Fine-tune Pipeline (Docker)") as demo:
    gr.Markdown("# Fine-tuning pipeline (LoRA) — Docker Space\nUsing Trendyol cybersecurity instruction dataset.")

    with gr.Tab("Train"):
        base_model = gr.Textbox(
            value="sshleifer/tiny-gpt2",
            label="Base model (HF Hub id)",
            info="Tip: for best chat behavior, use a small instruct/chat model and update LoRA target_modules in src/train.py accordingly.",
        )

        dataset_id = gr.Textbox(
            value="Trendyol/Trendyol-Cybersecurity-Instruction-Tuning-Dataset",
            label="Dataset (HF Hub id)",
            info="Expected columns: system, user, assistant (this dataset has them).",
        )

        with gr.Row():
            max_train_samples = gr.Number(value=2000, precision=0, label="Max train samples")
            max_steps = gr.Number(value=100, precision=0, label="Max steps")

        with gr.Row():
            lr = gr.Number(value=2e-4, label="Learning rate")
            batch_size = gr.Number(value=2, precision=0, label="Batch size")

        with gr.Row():
            lora_r = gr.Number(value=8, precision=0, label="LoRA r")
            lora_alpha = gr.Number(value=16, precision=0, label="LoRA alpha")
            lora_dropout = gr.Number(value=0.05, label="LoRA dropout")

        train_btn = gr.Button("Start fine-tune")
        train_out = gr.Textbox(lines=12, label="Status / logs")
        adapter_zip = gr.File(label="Download trained adapter (zip)")
        out_dir_box = gr.Textbox(label="Run output directory")
        adapter_dir_box = gr.Textbox(label="Adapter directory (use this in Generate tab)")

        train_btn.click(
            fn=run_train,
            inputs=[
                base_model,
                dataset_id,
                max_train_samples,
                max_steps,
                lr,
                batch_size,
                lora_r,
                lora_alpha,
                lora_dropout,
            ],
            outputs=[train_out, adapter_zip, out_dir_box, adapter_dir_box],
            queue=True,
        )

    with gr.Tab("Generate"):
        base_model2 = gr.Textbox(
            value="sshleifer/tiny-gpt2",
            label="Base model (must match training)",
        )
        adapter_dir = gr.Textbox(
            placeholder="Paste adapter dir path from Train tab (e.g., outputs/20260306-120000/adapter)",
            label="Adapter directory",
        )
        prompt = gr.Textbox(
            value="Explain the difference between phishing and spear phishing.",
            lines=4,
            label="Prompt",
        )

        with gr.Row():
            max_new_tokens = gr.Slider(16, 256, value=120, step=1, label="Max new tokens")
            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.05, label="Temperature")

        gen_btn = gr.Button("Generate")
        gen_out = gr.Textbox(lines=12, label="Output")

        gen_btn.click(
            fn=run_generate,
            inputs=[base_model2, adapter_dir, prompt, max_new_tokens, temperature],
            outputs=[gen_out],
            queue=False,
        )

demo.launch()
```