"""
Feature: Text Column Extract
============================

Extract **one text column** from a dataset (repo / subset / split) on the Hugging Face Hub
and write it as **one split** of another dataset.

Flow:
    1. Sign in (in the shared section of the app).
    2. Enter the source dataset: repo id, subset (config), split, text column name.
    3. Enter the destination dataset: repo id (free-form, NO hard-coded namespace) and split
       name (the "Generate split name" button derives a default from the source repo + subset).
    4. Click "Extract & push" — the work runs as a server-side job (see the Jobs tab).

All reads / writes go through the `datasets` + `huggingface_hub` libraries.
"""

from __future__ import annotations

import re
import traceback
from typing import Optional

import gradio as gr
from datasets import (
    concatenate_datasets,
    get_dataset_config_info,
    get_dataset_config_names,
    get_dataset_split_names,
    load_dataset,
)
from huggingface_hub import CommitOperationDelete, HfApi

from core.common import SPLIT_RE, AppContext, _clean, resolve_token, token_user
from core.jobs import submit_job

TITLE = "🧲 Text Column Extract"

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

PUSH_MODES = {
    "Overwrite the existing split": "overwrite",
    "Append to the existing split": "append",
    "Fail if the split already exists": "error",
}

MAX_SPLIT_NAME_LEN = 96


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def sanitize_split_name(name: str) -> str:
    """Turn any string into a split name accepted by the datasets library."""
    name = re.sub(r"[^0-9A-Za-z]+", "_", name or "")
    name = name.strip("_")[:MAX_SPLIT_NAME_LEN].strip("_")
    return name


def build_default_split_name(in_repo: str, in_subset: str, in_split: str, include_split: bool) -> str:
    """Default split name from the source repo + subset (+ the source split if requested)."""
    repo_name = _clean(in_repo).rstrip("/").split("/")[-1]
    parts = [repo_name]
    if _clean(in_subset):
        parts.append(_clean(in_subset))
    if include_split and _clean(in_split):
        parts.append(_clean(in_split))
    return sanitize_split_name("_".join(p for p in parts if p)) or "train"


def data_dir_for(config_name: str) -> str:
    """Parquet directory in the destination repo (matches `Dataset.push_to_hub`)."""
    return config_name if config_name != "default" else "data"


def shard_prefixes(data_dir: str, split: str) -> list[str]:
    """File prefixes that belong to exactly one split of a dataset repo."""
    return [f"{data_dir}/{split}-", f"{data_dir}/{split}.parquet"]


# ---------------------------------------------------------------------------
# "Generate split name" button
# ---------------------------------------------------------------------------


def generate_split_name(in_repo: str, in_subset: str, in_split: str, include_split: bool):
    name = build_default_split_name(in_repo, in_subset, in_split, include_split)
    if not SPLIT_RE.match(name):
        name = "train"
    return name, f"Generated split name: `{name}` — feel free to edit it."


# ---------------------------------------------------------------------------
# Inspect the source dataset
# ---------------------------------------------------------------------------


def inspect_source(
    oauth_token: Optional[gr.OAuthToken],
    session_token: Optional[str],
    in_repo: str,
    in_subset: str,
    in_column: str,
    in_revision: str,
):
    token = resolve_token(session_token, oauth_token)
    repo = _clean(in_repo)
    subset = _clean(in_subset) or None
    column = _clean(in_column)
    revision = _clean(in_revision) or None

    if not repo:
        return "❌ **Input repo** is empty."
    if not token:
        return "❌ Not signed in. Please sign in before inspecting."

    lines: list[str] = []
    try:
        configs = get_dataset_config_names(repo, token=token, revision=revision)
    except Exception as exc:  # noqa: BLE001
        return f"❌ Cannot read dataset `{repo}`.\n\n```\n{type(exc).__name__}: {exc}\n```"

    lines.append(f"### 🔎 `{repo}`")
    if len(configs) == 1 and configs[0] == "default":
        lines.append("- **Subset (config):** `default` — the dataset has a single subset.")
    else:
        shown = ", ".join(f"`{c}`" for c in configs[:25])
        if len(configs) > 25:
            shown += f" … (+{len(configs) - 25})"
        lines.append(f"- **Subsets:** {len(configs)} → {shown}")
        if subset and subset not in configs:
            lines.append(f"- ⚠️ Subset `{subset}` is **not** in the list above.")

    try:
        splits = get_dataset_split_names(repo, config_name=subset, token=token, revision=revision)
        lines.append(f"- **Splits of `{subset or 'default'}`:** " + ", ".join(f"`{s}`" for s in splits))
    except Exception as exc:  # noqa: BLE001
        lines.append(f"- ⚠️ Cannot read splits: `{type(exc).__name__}: {exc}`")

    try:
        info = get_dataset_config_info(repo, config_name=subset, token=token, revision=revision)
        columns = list(info.features.keys())
        lines.append(f"- **Columns:** {len(columns)}")
        lines.append("- **Column names:** " + ", ".join(f"`{c}`" for c in columns[:80]))
        if column:
            if column in columns:
                lines.append(f"- ✅ Column **`{column}`** exists.")
            else:
                lines.append(f"- ❌ Column **`{column}`** does NOT exist in this dataset.")
        sizes = [
            f"`{k}`: {v.num_examples:,}"
            for k, v in (getattr(info, "splits", None) or {}).items()
            if v.num_examples
        ]
        if sizes:
            lines.append("- **Rows:** " + ", ".join(sizes))
    except Exception as exc:  # noqa: BLE001
        lines.append(f"- ⚠️ Cannot read column info: `{type(exc).__name__}: {exc}`")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Extract & push to the Hub — a job generator (see `features/jobs.py`)
# ---------------------------------------------------------------------------


def _no_progress(*_args, **_kwargs) -> None:
    pass


def extract_and_push(
    token: str,
    in_repo: str,
    in_subset: str,
    in_split: str,
    in_column: str,
    in_revision: str,
    out_column: str,
    max_rows: float,
    num_proc: float,
    streaming: bool,
    out_repo: str,
    out_config: str,
    out_split: str,
    out_private: bool,
    push_mode_label: str,
    set_default: bool,
    max_shard_size: str,
    commit_message: str,
    progress=_no_progress,
):
    logs: list[str] = []

    def snapshot(status: str = "⏳ Working…"):
        return status, "\n".join(logs), None

    def log(message: str = "") -> tuple:
        logs.append(str(message))
        print(message, flush=True)
        return snapshot()

    def fail(message: str):
        log("")
        log(f"❌ ERROR: {message}")
        return f"### ❌ Failed\n\n{message}", "\n".join(logs), None

    try:
        # ---------------- 0. Token & input validation ----------------
        if not token:
            yield fail("Please **sign in** first (OAuth or Access Token).")
            return
        user = token_user(token)
        if not user:
            yield fail("Invalid or expired token.")
            return

        repo = _clean(in_repo)
        subset = _clean(in_subset) or None
        split = _clean(in_split)
        column = _clean(in_column)
        revision = _clean(in_revision) or None
        out_repo_id = _clean(out_repo)
        out_cfg = _clean(out_config) or "default"
        out_split_name = _clean(out_split)
        target_column = _clean(out_column) or column
        mode = PUSH_MODES.get(push_mode_label, "overwrite")
        n_rows = int(max_rows or 0)
        n_proc = int(num_proc or 0)
        shard_size = _clean(max_shard_size) or None

        if not repo:
            yield fail("**Input repo** is empty.")
            return
        if not split:
            yield fail("**Input split** is empty.")
            return
        if not column:
            yield fail("The **text column** to extract is empty.")
            return
        if not out_repo_id:
            yield fail("**Output repo** is empty (e.g. `your-name/dataset-name`).")
            return
        if not out_split_name:
            yield fail("**Output split** is empty. Click *Generate split name*.")
            return
        if not SPLIT_RE.match(out_split_name):
            yield fail(
                f"Split name `{out_split_name}` is invalid (only letters, digits, `_` and `.` "
                f"separators are allowed). Suggestion: `{sanitize_split_name(out_split_name)}`."
            )
            return
        if out_cfg == "data":
            yield fail("The config name cannot be `data`.")
            return
        if streaming and mode == "append":
            yield fail("**Streaming** mode cannot append to an existing split. Disable streaming.")
            return

        log(f"👤 Account: {user}")
        log("─" * 62)
        log("📥 SOURCE")
        log(f"   repo    : {repo}")
        log(f"   subset  : {subset or '(default)'}")
        log(f"   split   : {split}")
        log(f"   column  : {column}" + (f"  →  renamed to '{target_column}'" if target_column != column else ""))
        if revision:
            log(f"   revision: {revision}")
        log("📤 DESTINATION")
        log(f"   repo    : {out_repo_id}")
        log(f"   config  : {out_cfg}")
        log(f"   split   : {out_split_name}")
        log(f"   mode    : {mode}")
        log("─" * 62)
        yield snapshot("⏳ Loading the source dataset…")

        # ---------------- 1. Load the source dataset ----------------
        progress(0.05, desc="Loading the source dataset…")
        load_kwargs: dict = {
            "split": split,
            "token": token,
            "revision": revision,
            "streaming": bool(streaming),
        }
        if subset:
            load_kwargs["name"] = subset
        if n_proc > 1 and not streaming:
            load_kwargs["num_proc"] = n_proc

        source = load_dataset(repo, **load_kwargs)

        columns = list(source.column_names or [])
        if column not in columns:
            yield fail(
                f"Column `{column}` does not exist in `{repo}` / `{subset or 'default'}` / `{split}`.\n\n"
                f"Available columns: {', '.join(f'`{c}`' for c in columns[:60])}"
            )
            return

        # ---------------- 2. Extract the column ----------------
        progress(0.2, desc="Extracting the column…")
        yield log("⏳ Extracting the column…")
        extracted = source.select_columns([column])
        del source
        if target_column != column:
            extracted = extracted.rename_column(column, target_column)
            yield log(f"✏️ Renamed column `{column}` → `{target_column}`.")

        if n_rows > 0:
            if streaming:
                extracted = extracted.take(n_rows)
            else:
                extracted = extracted.select(range(min(n_rows, len(extracted))))
            yield log(f"🔢 Limited to at most {n_rows:,} rows.")

        if not streaming:
            yield log(f"✅ Extracted {len(extracted):,} rows (column `{target_column}`).")

        # ---------------- 3. Handle the destination repo / split ----------------
        api = HfApi(token=token)
        data_dir = data_dir_for(out_cfg)
        prefixes = shard_prefixes(data_dir, out_split_name)

        existing_files: list[str] = []
        repo_files: list[str] = []
        try:
            repo_files = list(api.list_repo_files(out_repo_id, repo_type="dataset", token=token))
            existing_files = [f for f in repo_files if any(f.startswith(p) for p in prefixes)]
            yield log(
                f"📦 Destination repo exists — {len(existing_files)} file(s) belong to split `{out_split_name}`."
            )
        except Exception as exc:  # noqa: BLE001
            yield log(f"ℹ️ Destination repo does not exist / is not accessible ({type(exc).__name__}). It will be created.")

        # `datasets` allows only one default config: if the repo already has a `default`
        # config (the `data/` directory), another config cannot be set as default.
        if out_cfg != "default" and set_default and any(f.startswith("data/") for f in repo_files):
            set_default = False
            yield log("ℹ️ The repo already has a `default` config → skipping `set_default` for the new config.")

        if existing_files and mode == "error":
            yield fail(
                f"Split `{out_split_name}` already exists in `{out_repo_id}` ({len(existing_files)} file(s)). "
                "Choose *Overwrite* or *Append*."
            )
            return

        if mode == "append" and existing_files:
            progress(0.35, desc="Reading the existing split to append…")
            yield log("⏳ Reading the existing split to append to it…")
            try:
                existing = load_dataset(
                    out_repo_id,
                    name=out_cfg if out_cfg != "default" else None,
                    split=out_split_name,
                    token=token,
                )
            except Exception as exc:  # noqa: BLE001
                yield fail(f"Cannot read the existing split `{out_split_name}` to append: {exc}")
                return

            if list(existing.column_names) != list(extracted.column_names):
                yield fail(
                    "Cannot append: the existing split and the new data have different columns.\n\n"
                    f"- Existing split : {existing.column_names}\n"
                    f"- New data       : {extracted.column_names}"
                )
                return
            yield log(f"➕ Existing split {len(existing):,} rows + new {len(extracted):,} rows.")
            extracted = concatenate_datasets([existing, extracted])
            del existing

        if existing_files and mode in ("overwrite", "append"):
            progress(0.45, desc="Deleting the split's old data…")
            yield log(f"🗑️ Deleting {len(existing_files)} old file(s) of split `{out_split_name}`…")
            api.create_commit(
                repo_id=out_repo_id,
                repo_type="dataset",
                operations=[CommitOperationDelete(path_in_repo=f) for f in existing_files],
                commit_message=f"Remove previous data of split '{out_split_name}'",
                token=token,
            )

        # `datasets` requires every split of the SAME config to share the same columns.
        if repo_files:
            try:
                out_info = get_dataset_config_info(
                    out_repo_id,
                    config_name=out_cfg if out_cfg != "default" else None,
                    token=token,
                )
                existing_columns = list(out_info.features.keys())
            except Exception:  # noqa: BLE001
                existing_columns = []
            if existing_columns and existing_columns != list(extracted.column_names):
                yield fail(
                    "Every split of the **same config** must have the same columns "
                    "(a constraint of the `datasets` library).\n\n"
                    f"- Config `{out_cfg}` has columns: `{existing_columns}`\n"
                    f"- The new data has columns: `{list(extracted.column_names)}`\n\n"
                    "→ Use a different **Output config**, or keep the column name (no rename)."
                )
                return

        # ---------------- 4. Push to the Hub ----------------
        progress(0.6, desc="Pushing to the Hugging Face Hub…")
        yield log("🚀 Pushing to the Hub (large datasets may take a few minutes)…")

        push_kwargs: dict = {
            "repo_id": out_repo_id,
            "config_name": out_cfg,
            "split": out_split_name,
            "private": bool(out_private),
            "token": token,
            "commit_message": _clean(commit_message)
            or f"Extract column '{column}' from {repo} ({subset or 'default'}/{split})",
        }
        if out_cfg != "default":
            push_kwargs["set_default"] = bool(set_default)
        if shard_size:
            push_kwargs["max_shard_size"] = shard_size
        if n_proc > 1:
            push_kwargs["num_proc"] = n_proc

        try:
            extracted.push_to_hub(**push_kwargs)
        except ValueError as exc:
            # Some `datasets` versions raise when a `default` config already exists.
            if push_kwargs.get("set_default") and "configuration named 'default'" in str(exc):
                yield log("⚠️ Cannot set_default (the repo already has a `default` config) → retrying without it…")
                push_kwargs.pop("set_default", None)
                extracted.push_to_hub(**push_kwargs)
            else:
                raise

        # ---------------- 5. Verify ----------------
        progress(0.92, desc="Verifying the result…")
        final_files = [
            f
            for f in api.list_repo_files(out_repo_id, repo_type="dataset", token=token)
            if any(f.startswith(p) for p in prefixes)
        ]
        total_bytes = 0
        try:
            for entry in api.list_repo_tree(out_repo_id, repo_type="dataset", token=token, recursive=True):
                path = getattr(entry, "path", "")
                size = getattr(entry, "size", None)
                if path in final_files and size:
                    total_bytes += size
        except Exception:  # noqa: BLE001
            pass

        rows_txt = "streaming" if streaming else f"{len(extracted):,} rows"
        url = f"https://huggingface.co/datasets/{out_repo_id}"
        log("")
        log(f"✅ DONE — {rows_txt}")
        for f in final_files:
            log(f"   • {f}")
        if total_bytes:
            log(f"   total size: {total_bytes / 1e6:.2f} MB")

        status = (
            "### ✅ Pushed successfully\n\n"
            f"- **Destination dataset:** [{out_repo_id}]({url})\n"
            f"- **Config / split:** `{out_cfg}` / `{out_split_name}`\n"
            f"- **Column:** `{target_column}` — {rows_txt}\n"
            f"- **Data files:** {len(final_files)}\n\n"
            "```python\n"
            "from datasets import load_dataset\n"
            f"ds = load_dataset({out_repo_id!r}"
            + (f", name={out_cfg!r}" if out_cfg != "default" else "")
            + f", split={out_split_name!r})\n"
            "```"
        )

        preview = None
        if not streaming:
            try:
                preview = extracted.select(range(min(5, len(extracted)))).to_pandas()
            except Exception:  # noqa: BLE001
                preview = None
        progress(1.0, desc="Done!")
        yield status, "\n".join(logs), preview

    except Exception as exc:  # noqa: BLE001
        log("")
        log(f"❌ {type(exc).__name__}: {exc}")
        log(traceback.format_exc())
        yield (
            f"### ❌ Failed\n\n```\n{type(exc).__name__}: {exc}\n```\n\n"
            "See the **Log** box for details.",
            "\n".join(logs),
            None,
        )


def submit_extract(
    oauth_token: Optional[gr.OAuthToken],
    session_token: Optional[str],
    in_repo: str,
    in_subset: str,
    in_split: str,
    in_column: str,
    in_revision: str,
    out_column: str,
    max_rows: float,
    num_proc: float,
    streaming: bool,
    out_repo: str,
    out_config: str,
    out_split: str,
    out_private: bool,
    push_mode_label: str,
    set_default: bool,
    max_shard_size: str,
    commit_message: str,
):
    """Queue `extract_and_push` as a server-side job; returns (message, job_id)."""
    params = {k: v for k, v in locals().items() if k not in ("oauth_token", "session_token")}
    source = "/".join(p for p in (_clean(in_repo), _clean(in_subset), _clean(in_split)) if p)
    title = f"{source} [{_clean(in_column)}] → {_clean(out_repo)}/{_clean(out_config) or 'default'}/{_clean(out_split)}"
    return submit_job(oauth_token, session_token, feature=TITLE, title=title, fn=extract_and_push, params=params)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------


def build(ctx: AppContext) -> None:
    """Build the tab inside the current `gr.Blocks`."""
    session_token = ctx.session_token
    gr.Markdown(
        "Extract **one text column** from a source dataset and write it as **one split** of a "
        "destination dataset on the Hugging Face Hub — using the `datasets` library."
    )

    # ---------------------- 2. Source ----------------------
    with gr.Accordion("📥 2. Source dataset", open=True):
        with gr.Row():
            in_repo = gr.Textbox(
                label="Input repo",
                placeholder="e.g. your-name/my-dataset  (or org/dataset-name)",
                scale=2,
            )
            in_subset = gr.Textbox(
                label="Subset (config)",
                placeholder="leave empty if none",
                scale=1,
            )
        with gr.Row():
            in_split = gr.Textbox(label="Input split", value="train", scale=1)
            in_column = gr.Textbox(label="Text column to extract", placeholder="e.g. text", scale=2)
        check_btn = gr.Button("🔎 Inspect source dataset")
        check_out = gr.Markdown()

        with gr.Accordion("⚙️ Advanced options (source)", open=False):
            with gr.Row():
                in_revision = gr.Textbox(label="Revision / branch (optional)", placeholder="main")
                out_column = gr.Textbox(
                    label="Rename the column in the output (optional)",
                    placeholder="empty = keep the name",
                )
            with gr.Row():
                max_rows = gr.Number(label="Row limit (0 = all)", value=0, precision=0)
                num_proc = gr.Number(label="num_proc (0/1 = default)", value=0, precision=0)
                streaming = gr.Checkbox(
                    label="Streaming (very large datasets)",
                    value=False,
                    info="Does not load everything into RAM; append mode is not supported.",
                )

    # ---------------------- 3. Destination ----------------------
    with gr.Accordion("📤 3. Destination dataset", open=True):
        out_repo = gr.Textbox(
            label="Output repo",
            placeholder="e.g. your-name/new-dataset  (free-form — any namespace)",
            info="Your repo or an organization's repo you can write to. Created if it does not exist.",
        )
        with gr.Row():
            out_config = gr.Textbox(label="Output config (subset)", value="default", scale=1)
            out_split = gr.Textbox(label="Output split", placeholder="click Generate split name →", scale=2)
        with gr.Row():
            gen_btn = gr.Button("✨ Generate split name", variant="secondary")
            use_input_repo_btn = gr.Button("↩️ Reuse Input repo", variant="secondary")
            include_split_in_name = gr.Checkbox(
                label="Also include the source split name", value=False, scale=1
            )
        gen_hint = gr.Markdown(
            "Split names may only contain letters, digits, `_` (and `.` separators). Example: `my_dataset_vi_v2`."
        )

        with gr.Accordion("⚙️ Advanced options (destination)", open=False):
            out_private = gr.Checkbox(label="Create the repo as private", value=False)
            push_mode = gr.Radio(
                choices=list(PUSH_MODES.keys()),
                value="Overwrite the existing split",
                label="If the split already exists",
            )
            set_default = gr.Checkbox(
                label="Make this config the default one (set_default)",
                value=True,
                info="Only applies when Output config is not 'default'.",
            )
            max_shard_size = gr.Textbox(label="max_shard_size", value="500MB")
            commit_message = gr.Textbox(
                label="Commit message (optional)", placeholder="empty = auto-generated"
            )

    # ---------------------- 4. Run ----------------------
    extract_btn = gr.Button("🚀 4. Submit job: extract & push to the Hub", variant="primary", elem_classes=["run-btn"])
    submit_status = gr.Markdown()

    gr.Markdown(
        """
        ---
        **Notes**
        - Data is written as Parquet: `data/<split>-00000-of-00001.parquet`
          (or `<config>/<split>-....parquet` when Output config is not `default`).
        - *Overwrite* / *Append* **delete the old files of that split only** before pushing → no duplicates.
        - Every split of the **same config** must have the **same columns** (a `datasets` constraint).
          To rename the column, use a different Output config.
        - The job runs on the server: follow its log and a preview of the first 5 rows in the **📋 Jobs** tab.
        - The token is kept only in your browser session and in memory while its job runs, never written to disk.
        """
    )

    # ---------------------- Events ----------------------
    # Note: `gr.OAuthToken` parameters are injected by Gradio from the type hint and
    # must NOT be listed in `inputs`.
    check_btn.click(
        inspect_source,
        inputs=[session_token, in_repo, in_subset, in_column, in_revision],
        outputs=[check_out],
    )
    gen_btn.click(
        generate_split_name,
        inputs=[in_repo, in_subset, in_split, include_split_in_name],
        outputs=[out_split, gen_hint],
    )
    use_input_repo_btn.click(lambda repo: repo, inputs=[in_repo], outputs=[out_repo])
    extract_btn.click(
        submit_extract,
        inputs=[
            session_token,
            in_repo,
            in_subset,
            in_split,
            in_column,
            in_revision,
            out_column,
            max_rows,
            num_proc,
            streaming,
            out_repo,
            out_config,
            out_split,
            out_private,
            push_mode,
            set_default,
            max_shard_size,
            commit_message,
        ],
        outputs=[submit_status, ctx.submitted_job],
        concurrency_limit=None,
    )