"""Minimal ROBERT web interface for Hugging Face Spaces."""

from __future__ import annotations

import spaces
import os
import queue
import shutil
import signal
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path
from typing import Generator

import gradio as gr
import pandas as pd


APP_TITLE = "ROBERT"
JOB_ROOT = Path(tempfile.gettempdir()) / "robert_web_jobs"
JOB_ROOT.mkdir(parents=True, exist_ok=True)

# Conservative defaults for a public, shared Space. They can be changed in the
# Space Settings by adding variables with these names.
MAX_UPLOAD_MB = int(os.getenv("ROBERT_MAX_UPLOAD_MB", "20"))
MAX_ROWS = int(os.getenv("ROBERT_MAX_ROWS", "2000"))
# Free accounts receive five minutes of ZeroGPU usage per day. ROBERT is
# stopped a little earlier so there is time to collect and return its files.
ZERO_GPU_DURATION_SECONDS = 120
TIMEOUT_SECONDS = min(
    int(os.getenv("ROBERT_TIMEOUT_SECONDS", "285")),
    ZERO_GPU_DURATION_SECONDS - 15,
)
JOB_TTL_SECONDS = int(os.getenv("ROBERT_JOB_TTL_SECONDS", "21600"))
MAX_LOG_CHARS = 120_000
MAX_GALLERY_IMAGES = 24


CSS = """
.gradio-container {
    max-width: 1120px !important;
    margin: 0 auto !important;
}
.hero {
    padding: 2.4rem 1.2rem 1.6rem;
    text-align: center;
}
.hero .eyebrow {
    color: #4f46e5;
    font-size: .78rem;
    font-weight: 700;
    letter-spacing: .14em;
    text-transform: uppercase;
}
.hero h1 {
    font-size: clamp(2.8rem, 8vw, 5.2rem);
    line-height: .95;
    margin: .45rem 0 .8rem;
    letter-spacing: -.06em;
}
.hero p {
    color: #64748b;
    font-size: 1.05rem;
    margin: 0 auto;
    max-width: 700px;
}
.section-title {
    margin-top: .35rem;
}
#run-button {
    background: linear-gradient(90deg, #4f46e5, #2563eb) !important;
    border: 0 !important;
}
.footer {
    color: #64748b;
    font-size: .86rem;
    padding: 1.8rem .4rem .6rem;
    text-align: center;
}
"""


def _empty_column_controls() -> tuple[gr.Dropdown, gr.Dropdown, gr.Dropdown]:
    return (
        gr.Dropdown(choices=[], value=None),
        gr.Dropdown(choices=[], value=None),
        gr.Dropdown(choices=[], value=[]),
    )


def _read_csv(path: str | Path, *, preview_only: bool = False) -> pd.DataFrame:
    """Read common CSV variants and fail with a concise user-facing message."""
    source = Path(path)
    if not source.is_file():
        raise ValueError("The uploaded file is no longer available.")

    size_mb = source.stat().st_size / (1024 * 1024)
    if size_mb > MAX_UPLOAD_MB:
        raise ValueError(
            f"The file is {size_mb:.1f} MB. The current limit is {MAX_UPLOAD_MB} MB."
        )

    kwargs: dict[str, object] = {
        "sep": None,
        "engine": "python",
        "encoding": "utf-8-sig",
    }
    if preview_only:
        kwargs["nrows"] = 100

    try:
        dataframe = pd.read_csv(source, **kwargs)
    except UnicodeDecodeError:
        kwargs["encoding"] = "latin-1"
        dataframe = pd.read_csv(source, **kwargs)
    except Exception as exc:
        raise ValueError(f"The CSV could not be read: {exc}") from exc

    if dataframe.empty:
        raise ValueError("The CSV is empty.")
    if len(dataframe.columns) < 3:
        raise ValueError(
            "The CSV must contain at least a name/ID column, a target column "
            "and one descriptor column."
        )
    return dataframe


def _guess_column(columns: list[str], candidates: tuple[str, ...]) -> str | None:
    lowered = {column.strip().lower(): column for column in columns}
    for candidate in candidates:
        if candidate in lowered:
            return lowered[candidate]
    return None


def inspect_training_csv(
    csv_path: str | None,
) -> tuple[pd.DataFrame, gr.Dropdown, gr.Dropdown, gr.Dropdown, str]:
    """Preview a dataset and populate all column selectors."""
    if not csv_path:
        target, names, ignored = _empty_column_controls()
        return pd.DataFrame(), target, names, ignored, "Upload a CSV to begin."

    try:
        dataframe = _read_csv(csv_path, preview_only=True)
    except ValueError as exc:
        target, names, ignored = _empty_column_controls()
        return pd.DataFrame(), target, names, ignored, f"❌ {exc}"

    columns = [str(column) for column in dataframe.columns]
    names_value = _guess_column(
        columns, ("name", "names", "id", "entry", "code", "compound")
    )
    target_value = _guess_column(
        columns,
        ("target", "target_values", "y", "response", "property", "activity"),
    )

    if names_value is None:
        names_value = columns[0]
    if target_value is None:
        target_value = next(
            (column for column in columns if column != names_value), columns[0]
        )

    return (
        dataframe,
        gr.Dropdown(choices=columns, value=target_value),
        gr.Dropdown(choices=columns, value=names_value),
        gr.Dropdown(choices=columns, value=[]),
        f"✅ CSV detected: {len(columns)} columns. Showing the first "
        f"{len(dataframe)} rows.",
    )


def _normalize_csv(source_path: str, destination: Path) -> pd.DataFrame:
    dataframe = _read_csv(source_path)
    if len(dataframe) > MAX_ROWS:
        raise ValueError(
            f"The dataset contains {len(dataframe):,} rows. The public Space is "
            f"limited to {MAX_ROWS:,} rows."
        )
    dataframe.to_csv(destination, index=False)
    return dataframe


def _cleanup_old_jobs() -> None:
    """Remove stale temporary jobs while keeping recent downloads available."""
    cutoff = time.time() - JOB_TTL_SECONDS
    for item in JOB_ROOT.iterdir():
        try:
            if item.stat().st_mtime >= cutoff:
                continue
            if item.is_dir():
                shutil.rmtree(item)
            elif item.is_file():
                item.unlink()
        except OSError:
            # Cleanup must never prevent a new calculation from starting.
            continue


def _terminate_process(process: subprocess.Popen[str]) -> None:
    if process.poll() is not None:
        return
    try:
        os.killpg(process.pid, signal.SIGTERM)
        process.wait(timeout=5)
    except (ProcessLookupError, subprocess.TimeoutExpired):
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


def _build_command(
    training_name: str,
    target: str,
    names: str,
    problem_type: str,
    ignored_columns: list[str] | None,
    test_name: str | None,
) -> list[str]:
    command = [
        sys.executable,
        "-u",
        "-m",
        "robert",
        "--csv_name",
        training_name,
        "--y",
        target,
        "--names",
        names,
    ]

    if problem_type == "Classification":
        command.extend(["--type", "clas"])
    if ignored_columns:
        command.extend(["--ignore", repr(ignored_columns)])
    if test_name:
        command.extend(["--csv_test", test_name])
    return command


def _collect_figures(run_directory: Path) -> list[tuple[str, str]]:
    extensions = {".png", ".jpg", ".jpeg"}
    paths = sorted(
        path
        for path in run_directory.rglob("*")
        if path.is_file() and path.suffix.lower() in extensions
    )[:MAX_GALLERY_IMAGES]
    return [(str(path), path.name) for path in paths]


@spaces.GPU(duration=ZERO_GPU_DURATION_SECONDS)
def run_robert(
    training_csv: str | None,
    target: str | None,
    names: str | None,
    problem_type: str,
    ignored_columns: list[str] | None,
    test_csv: str | None,
) -> Generator[tuple[str, str, str | None, str | None, list[tuple[str, str]]], None, None]:
    """Run ROBERT in an isolated directory and stream its console output."""
    process: subprocess.Popen[str] | None = None
    log_lines: list[str] = []

    try:
        if not training_csv:
            raise ValueError("Upload a training CSV before running ROBERT.")
        if not target or not names:
            raise ValueError("Select both the target and name/ID columns.")
        if target == names:
            raise ValueError("Target and name/ID must be different columns.")

        ignored_columns = ignored_columns or []
        forbidden_ignored = sorted({target, names}.intersection(ignored_columns))
        if forbidden_ignored:
            raise ValueError(
                "Target and name/ID cannot be ignored: " + ", ".join(forbidden_ignored)
            )

        _cleanup_old_jobs()
        run_directory = Path(tempfile.mkdtemp(prefix="job_", dir=JOB_ROOT))
        training_path = run_directory / "training.csv"
        training_data = _normalize_csv(training_csv, training_path)

        missing = [column for column in (target, names) if column not in training_data.columns]
        if missing:
            raise ValueError("Columns not found in the CSV: " + ", ".join(missing))

        test_name: str | None = None
        if test_csv:
            test_path = run_directory / "external_test.csv"
            _normalize_csv(test_csv, test_path)
            test_name = test_path.name

        command = _build_command(
            training_path.name,
            target,
            names,
            problem_type,
            ignored_columns,
            test_name,
        )

        environment = os.environ.copy()
        environment.update(
            {
                "PYTHONUNBUFFERED": "1",
                "MPLBACKEND": "Agg",
                "MPLCONFIGDIR": str(run_directory / ".matplotlib"),
                "OMP_NUM_THREADS": "2",
                "OPENBLAS_NUM_THREADS": "2",
                "MKL_NUM_THREADS": "2",
                "NUMBA_NUM_THREADS": "2",
                "QT_QPA_PLATFORM": "offscreen",
            }
        )

        process = subprocess.Popen(
            command,
            cwd=run_directory,
            env=environment,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
            bufsize=1,
            start_new_session=True,
        )

        output_queue: queue.Queue[str] = queue.Queue()

        def read_output() -> None:
            assert process is not None and process.stdout is not None
            for output_line in process.stdout:
                output_queue.put(output_line.rstrip())

        reader = threading.Thread(target=read_output, daemon=True)
        reader.start()
        started_at = time.monotonic()
        last_update = 0.0

        log_lines.append("ROBERT calculation started.")
        yield "⏳ ROBERT is running…", "\n".join(log_lines), None, None, []

        timed_out = False
        while process.poll() is None or not output_queue.empty():
            received_output = False
            while True:
                try:
                    log_lines.append(output_queue.get_nowait())
                    received_output = True
                except queue.Empty:
                    break

            if time.monotonic() - started_at > TIMEOUT_SECONDS:
                timed_out = True
                _terminate_process(process)
                log_lines.append(
                    f"ERROR: calculation exceeded the {TIMEOUT_SECONDS // 60}-minute limit."
                )
                break

            now = time.monotonic()
            if received_output and now - last_update >= 0.35:
                visible_log = "\n".join(log_lines)[-MAX_LOG_CHARS:]
                yield "⏳ ROBERT is running…", visible_log, None, None, []
                last_update = now
            time.sleep(0.12)

        reader.join(timeout=1)
        while not output_queue.empty():
            log_lines.append(output_queue.get_nowait())

        visible_log = "\n".join(log_lines)[-MAX_LOG_CHARS:]
        if timed_out:
            yield "❌ The calculation timed out.", visible_log, None, None, []
            return
        if process.returncode != 0:
            yield (
                f"❌ ROBERT stopped with exit code {process.returncode}.",
                visible_log,
                None,
                None,
                [],
            )
            return

        reports = sorted(run_directory.rglob("ROBERT_report.pdf"))
        report_path = str(reports[0]) if reports else None
        figures = _collect_figures(run_directory)

        archive_base = JOB_ROOT / f"{run_directory.name}_ROBERT_results"
        archive_path = shutil.make_archive(
            str(archive_base), "zip", root_dir=run_directory
        )

        status = "✅ ROBERT finished successfully."
        if report_path is None:
            status += " No ROBERT_report.pdf was found; check the log and ZIP archive."
        yield status, visible_log, report_path, archive_path, figures

    except Exception as exc:
        log_lines.append(f"ERROR: {exc}")
        yield "❌ The calculation could not be started.", "\n".join(log_lines), None, None, []
    finally:
        if process is not None and process.poll() is None:
            _terminate_process(process)


with gr.Blocks(title=APP_TITLE) as demo:
    gr.HTML(
        """
        <div class="hero">
          <div class="eyebrow">Automated machine learning for chemistry</div>
          <h1>ROBERT</h1>
          <p>
            Upload a dataset, select the response and let ROBERT curate the data,
            build models, validate them and prepare a publication-ready report.
          </p>
        </div>
        """
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            gr.Markdown("### 1 · Dataset", elem_classes="section-title")
            training_file = gr.File(
                label="Training dataset (.csv)",
                file_types=[".csv"],
                type="filepath",
            )
            file_status = gr.Markdown("Upload a CSV to begin.")
            preview = gr.Dataframe(
                label="Dataset preview",
                interactive=False,
                wrap=True,
                max_height=330,
            )

        with gr.Column(scale=4):
            gr.Markdown("### 2 · Configuration", elem_classes="section-title")
            target_column = gr.Dropdown(label="Target column", choices=[])
            names_column = gr.Dropdown(label="Name / ID column", choices=[])
            problem_type = gr.Radio(
                label="Problem type",
                choices=["Regression", "Classification"],
                value="Regression",
            )

            with gr.Accordion("Optional settings", open=False):
                ignored_columns = gr.Dropdown(
                    label="Columns to ignore",
                    choices=[],
                    multiselect=True,
                    value=[],
                )
                test_file = gr.File(
                    label="External test dataset (.csv)",
                    file_types=[".csv"],
                    type="filepath",
                )

            with gr.Row():
                run_button = gr.Button(
                    "Run ROBERT", variant="primary", elem_id="run-button", scale=3
                )
                stop_button = gr.Button("Stop", variant="stop", scale=1)

    gr.Markdown("### 3 · Results", elem_classes="section-title")
    run_status = gr.Markdown("Results will appear here.")
    with gr.Tabs():
        with gr.Tab("Report & downloads"):
            with gr.Row():
                report_file = gr.File(label="ROBERT report (PDF)", interactive=False)
                results_archive = gr.File(
                    label="All ROBERT results (ZIP)", interactive=False
                )
        with gr.Tab("Figures"):
            figures_gallery = gr.Gallery(
                label="Generated figures",
                columns=3,
                object_fit="contain",
            )
        with gr.Tab("Live log"):
            console = gr.Textbox(
                label="ROBERT output",
                lines=20,
                max_lines=30,
                interactive=False,
            )

    gr.HTML(
        """
        <div class="footer">
          ROBERT ·
          <a href="https://github.com/jvalegre/robert/tree/ddg_branch" target="_blank">GitHub</a>
          ·
          <a href="https://robert.readthedocs.io/" target="_blank">Documentation</a>
          <br>Uploaded files and results are temporary. Do not upload confidential data.
        </div>
        """
    )

    training_file.upload(
        fn=inspect_training_csv,
        inputs=training_file,
        outputs=[
            preview,
            target_column,
            names_column,
            ignored_columns,
            file_status,
        ],
    )
    training_file.clear(
        fn=inspect_training_csv,
        inputs=training_file,
        outputs=[
            preview,
            target_column,
            names_column,
            ignored_columns,
            file_status,
        ],
    )

    run_event = run_button.click(
        fn=run_robert,
        inputs=[
            training_file,
            target_column,
            names_column,
            problem_type,
            ignored_columns,
            test_file,
        ],
        outputs=[
            run_status,
            console,
            report_file,
            results_archive,
            figures_gallery,
        ],
        concurrency_limit=1,
        concurrency_id="robert_calculation",
    )
    stop_button.click(fn=None, cancels=[run_event])


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=5).launch(
        theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
        css=CSS,
    )