from __future__ import annotations

import html
import json
from functools import lru_cache

import gradio as gr
import pandas as pd
from datasets import load_dataset


DATASET_ID = "suyuan37/TS-Dataset"
CONFIG_FIELDS = {
    "sft": {
        "answer_fields": ("input", "output"),
        "labels": ["全部", "accident-description", "causal-analysis", "traffic-knowledge", "general-knowledge"],
    },
    "dpo": {
        "answer_fields": ("chosen", "rejected"),
        "labels": ["全部"],
    },
}


@lru_cache(maxsize=4)
def get_frame(config_name: str) -> pd.DataFrame:
    dataset = load_dataset(DATASET_ID, config_name, split="train")
    frame = dataset.to_pandas()
    frame.insert(0, "row", range(len(frame)))
    return frame


RESULT_COLUMNS = ["row", "label", "instruction", "preview"]


def render_json_markdown(value: object, level: int = 0) -> str:
    prefix = "  " * level
    if isinstance(value, dict):
        lines = []
        for key, item in value.items():
            safe_key = html.escape(str(key), quote=False)
            if isinstance(item, (dict, list)):
                lines.append(f"{prefix}- **{safe_key}**")
                lines.append(render_json_markdown(item, level + 1))
            else:
                safe_item = html.escape(str(item), quote=False)
                lines.append(f"{prefix}- **{safe_key}**：{safe_item}")
        return "\n".join(lines)
    if isinstance(value, list):
        lines = []
        for item in value:
            if isinstance(item, (dict, list)):
                lines.append(f"{prefix}-")
                lines.append(render_json_markdown(item, level + 1))
            else:
                lines.append(f"{prefix}- {html.escape(str(item), quote=False)}")
        return "\n".join(lines)
    return f"{prefix}{html.escape(str(value), quote=False)}"


def as_markdown(value: object, *, json_block: bool = False) -> str:
    if value is None:
        return "_空_"
    text = str(value).replace("\r\n", "\n").replace("\r", "\n").strip()
    if not text:
        return "_空_"
    if json_block:
        try:
            parsed = json.loads(text)
        except json.JSONDecodeError:
            return f"```json\n{text}\n```"
        return render_json_markdown(parsed)
    return html.escape(text, quote=False)


def compact(value: object, limit: int = 140) -> str:
    text = "" if value is None else str(value)
    text = " ".join(text.replace("\r", "\n").split())
    return text if len(text) <= limit else text[: limit - 1] + "..."


def label_choices(config_name: str) -> list[str]:
    return CONFIG_FIELDS[config_name]["labels"]


def render_row(config_name: str, row_number: int):
    frame = get_frame(config_name)
    if frame.empty:
        return (
            "未找到数据。",
            "_空_",
            gr.update(value="_空_", visible=config_name == "sft"),
            gr.update(value="_空_", visible=config_name == "sft"),
            gr.update(value="_空_", visible=config_name == "dpo"),
            gr.update(value="_空_", visible=config_name == "dpo"),
        )

    row_number = max(0, min(int(row_number), len(frame) - 1))
    row = frame.iloc[row_number].to_dict()
    label = row.get("label", "dpo-pair")
    meta = f"**{config_name.upper()}** · row `{row_number}` / `{len(frame) - 1}` · `{label}`"

    instruction = as_markdown(row.get("instruction"))
    input_md = as_markdown(row.get("input"))
    output_md = as_markdown(row.get("output"))
    chosen_md = as_markdown(row.get("chosen"), json_block=config_name == "dpo")
    rejected_md = as_markdown(row.get("rejected"), json_block=config_name == "dpo")

    return (
        meta,
        instruction,
        gr.update(value=input_md, visible=config_name == "sft"),
        gr.update(value=output_md, visible=config_name == "sft"),
        gr.update(value=chosen_md, visible=config_name == "dpo"),
        gr.update(value=rejected_md, visible=config_name == "dpo"),
    )


def configure_dataset(config_name: str):
    frame = get_frame(config_name)
    max_row = max(len(frame) - 1, 0)
    table = search_rows(config_name, "", "全部", 12)
    rendered = render_row(config_name, 0)
    return (
        gr.update(maximum=max_row, value=0),
        gr.update(choices=label_choices(config_name), value="全部", visible=config_name == "sft"),
        table,
        *rendered,
    )


def search_rows(config_name: str, query: str, label_filter: str, limit: int) -> pd.DataFrame:
    frame = get_frame(config_name)
    limit = int(limit or 12)
    query = (query or "").strip().lower()

    view = frame
    if config_name == "sft" and label_filter and label_filter != "全部":
        view = view[view["label"] == label_filter]

    if query:
        fields = ["instruction", *CONFIG_FIELDS[config_name]["answer_fields"]]

        def matched(row: pd.Series) -> bool:
            return any(query in str(row.get(field, "")).lower() for field in fields)

        view = view[view.apply(matched, axis=1)]

    rows = []
    for _, row in view.head(limit).iterrows():
        if config_name == "sft":
            rows.append([int(row["row"]), row["label"], compact(row["instruction"], 120), compact(row["output"], 180)])
        else:
            rows.append(
                [
                    int(row["row"]),
                    "dpo-pair",
                    compact(row["instruction"], 120),
                    f"chosen: {compact(row['chosen'], 120)}\nrejected: {compact(row['rejected'], 120)}",
                ]
            )

    return pd.DataFrame(rows, columns=RESULT_COLUMNS)


def step_row(config_name: str, row_number: int, step: int):
    frame = get_frame(config_name)
    next_row = max(0, min(int(row_number) + step, len(frame) - 1))
    return (gr.update(value=next_row), *render_row(config_name, next_row))


def pick_result(config_name: str, table_value, evt: gr.SelectData):
    try:
        table = pd.DataFrame(table_value, columns=RESULT_COLUMNS)
        picked_row = int(table.iloc[int(evt.index[0])]["row"])
    except Exception:
        picked_row = 0
    return (gr.update(value=picked_row), *render_row(config_name, picked_row))


css = """
:root {
  --sr-ink: #17211f;
  --sr-muted: #5d6b67;
  --sr-line: #cad8d3;
  --sr-panel: #f7fbf9;
  --sr-accent: #0f8a7b;
}
.gradio-container {
  background: linear-gradient(180deg, #eef8f5 0%, #ffffff 24%, #ffffff 100%);
}
.sr-title {
  border-left: 6px solid var(--sr-accent);
  padding: 0.4rem 0 0.4rem 1rem;
}
.sr-title h1 {
  margin: 0;
  color: var(--sr-ink);
}
.sr-title p {
  margin: 0.25rem 0 0;
  color: var(--sr-muted);
}
.sr-panel {
  border: 1px solid var(--sr-line);
  border-radius: 8px;
  background: var(--sr-panel);
  padding: 0.9rem 1rem;
}
.sr-panel h2, .sr-panel h3 {
  margin-top: 0.2rem;
}
"""


with gr.Blocks(css=css, title="SafetyRAISE TS Dataset Viewer") as demo:
    gr.HTML(
        """
        <div class="sr-title">
          <h1>SafetyRAISE TS Dataset Viewer</h1>
          <p>Markdown preview for SFT instructions, accident contexts, model answers, and DPO preference pairs.</p>
        </div>
        """
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=260):
            config = gr.Dropdown(["sft", "dpo"], value="sft", label="Subset")
            row_number = gr.Slider(0, maximum=len(get_frame("sft")) - 1, value=0, step=1, label="Row")
            with gr.Row():
                previous_btn = gr.Button("Previous")
                next_btn = gr.Button("Next")
            label_filter = gr.Dropdown(label_choices("sft"), value="全部", label="Label")
            query = gr.Textbox(label="Search", placeholder="事故、责任、道路类型...")
            result_limit = gr.Slider(5, 50, value=12, step=1, label="Results")
            search_btn = gr.Button("Search", variant="primary")
        with gr.Column(scale=2):
            meta = gr.Markdown(elem_classes=["sr-panel"])
            instruction_md = gr.Markdown(elem_classes=["sr-panel"])

    with gr.Row():
        input_md = gr.Markdown(label="Input", elem_classes=["sr-panel"])
        output_md = gr.Markdown(label="Output", elem_classes=["sr-panel"])

    with gr.Row():
        chosen_md = gr.Markdown(label="Chosen", visible=False, elem_classes=["sr-panel"])
        rejected_md = gr.Markdown(label="Rejected", visible=False, elem_classes=["sr-panel"])

    results = gr.Dataframe(
        headers=RESULT_COLUMNS,
        datatype=["number", "str", "str", "str"],
        interactive=False,
        label="Search results",
        wrap=True,
    )

    outputs = [meta, instruction_md, input_md, output_md, chosen_md, rejected_md]
    config.change(configure_dataset, inputs=[config], outputs=[row_number, label_filter, results, *outputs])
    row_number.change(render_row, inputs=[config, row_number], outputs=outputs)
    previous_btn.click(lambda cfg, row: step_row(cfg, row, -1), inputs=[config, row_number], outputs=[row_number, *outputs])
    next_btn.click(lambda cfg, row: step_row(cfg, row, 1), inputs=[config, row_number], outputs=[row_number, *outputs])
    search_btn.click(search_rows, inputs=[config, query, label_filter, result_limit], outputs=[results])
    results.select(pick_result, inputs=[config, results], outputs=[row_number, *outputs])

    demo.load(configure_dataset, inputs=[config], outputs=[row_number, label_filter, results, *outputs])


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