#!/usr/bin/env python3
"""omniopt_evaluate - Graphical/textual browse for OmniOpt runs.

Python rewrite of the former ``omniopt_evaluate`` bash script.  The
original used ``whiptail`` for every interactive prompt; this version
keeps the same logical flow but replaces every dialog with a textual
prompt (printed menu + stdin input).  Behavioural compatibility:

  * same CLI flags (``--projectdir``, ``--debug``, ``--nogauge``,
    ``--dont_load_modules``, ``--no_upgrade``, ``--help``)
  * same project-discovery logic
  * same wallclock / failed-job counting
  * same error-handling exit codes
"""

from __future__ import annotations

import argparse
import csv
import os
import re
import shutil
import sys
from pathlib import Path
from typing import Callable, Iterable, List, Optional, Sequence, Tuple


SCRIPT_DIR = Path(__file__).resolve().parent


# ---------------------------------------------------------------------------
# Pure data helpers (used by tests)
# ---------------------------------------------------------------------------


def extract_experiment_name(run_dir: str) -> str:
    """Extract experiment name from a path like ``runs/myexp/0``.

    Mirrors the bash version: strip a trailing ``/[0-9]+`` and then
    take the last component.
    """
    s = run_dir.rstrip("/")
    s = re.sub(r"/[0-9]+$", "", s)
    return os.path.basename(s)


def calculate_wallclock_time(csv_path: str) -> int:
    """Read ``start_time,end_time`` from ``csv_path`` and return the
    elapsed time in whole seconds (matches the bash integer division).
    Returns 0 on missing or empty files.
    """
    try:
        with open(csv_path, "r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            starts: List[float] = []
            ends: List[float] = []
            for row in reader:
                try:
                    starts.append(float(row["start_time"]))
                    ends.append(float(row["end_time"]))
                except (KeyError, ValueError):
                    continue
    except FileNotFoundError:
        return 0

    if not starts:
        return 0

    min_t = int(min(starts))
    max_t = int(max(ends))
    return max(0, max_t - min_t)


def count_failed_jobs(csv_path: str) -> int:
    """Count rows where ``exit_code`` is non-zero. Returns 0 on missing file."""
    try:
        with open(csv_path, "r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            n = 0
            for row in reader:
                try:
                    if int(row.get("exit_code", "0")) != 0:
                        n += 1
                except ValueError:
                    continue
            return n
    except FileNotFoundError:
        return 0


def format_duration(seconds: int) -> str:
    """Format a duration in seconds as ``Nd Nh Nm Ns`` (matches bash)."""
    seconds = int(seconds)
    days = seconds // 86400
    hours = (seconds % 86400) // 3600
    minutes = (seconds % 3600) // 60
    secs = seconds % 60
    parts: List[str] = []
    if days:
        parts.append(f"{days}d")
    if hours:
        parts.append(f"{hours}h")
    if minutes:
        parts.append(f"{minutes}m")
    parts.append(f"{secs}s")
    return " ".join(parts)


def find_projects(projectdir: str) -> List[str]:
    """List experiment names under ``projectdir`` that have a results.csv
    somewhere in their run sub-directories.

    Mirrors the bash logic: ``ls $PROJECTDIR/*/*/results.csv`` and the
    parent directory name.
    """
    root = Path(projectdir)
    if not root.is_dir():
        return []
    found: set[str] = set()
    for results_csv in root.glob("*/*/results.csv"):
        # /root/<project>/<run>/results.csv  ->  <project>
        found.add(results_csv.parent.parent.name)
    return sorted(found)


def find_run_numbers(project_dir: str) -> List[str]:
    """List run numbers (subdirectory names that are all digits) under
    ``project_dir`` that have a results.csv inside them."""
    root = Path(project_dir)
    if not root.is_dir():
        return []
    out: List[str] = []
    for child in sorted(root.iterdir()):
        if not child.is_dir() or not child.name.isdigit():
            continue
        if (child / "results.csv").exists():
            out.append(child.name)
    return out


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="omniopt_evaluate",
        description="Browse OmniOpt runs (textual interface).",
        add_help=False,
    )
    p.add_argument("--help", action="store_true", help="Show help.")
    p.add_argument("--debug", action="store_true", help="Enable debug output.")
    p.add_argument("--nogauge", action="store_true", help="Disable gauges.")
    p.add_argument(
        "--projectdir", default="runs", help="Path to projects (default: runs)."
    )
    p.add_argument(
        "--dont_load_modules",
        action="store_true",
        help="Don't load modules (no-op in Python version).",
    )
    p.add_argument(
        "--no_upgrade", action="store_true", help="Disable upgrades (no-op)."
    )
    return p


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    return _build_parser().parse_args(list(argv))


def parse_args_with_exit_code(argv: Sequence[str]) -> int:
    try:
        parse_args(argv)
    except SystemExit as e:
        return int(e.code) if isinstance(e.code, int) else 1
    return 0


# ---------------------------------------------------------------------------
# Textual prompts (replace whiptail)
# ---------------------------------------------------------------------------


def textual_input(
    prompt: str,
    *,
    default: str = "",
    input_fn: Callable[[str], str] = input,
    output_fn: Callable[[str], None] = print,
) -> str:
    """Read one line from ``input_fn``; fall back to ``default`` if empty."""
    suffix = f" [{default}]" if default else ""
    output_fn(f"{prompt}{suffix}: ")
    line = input_fn("").strip()
    return line or default


def textual_menu(
    title: str,
    options: Sequence[Tuple[str, str]],
    *,
    input_fn: Callable[[str], str] = input,
    output_fn: Callable[[str], None] = print,
) -> str:
    """Print a numbered menu and return the chosen key.

    ``options`` is a list of ``(key, description)`` tuples.  Accepts
    either the key (e.g. ``q``) or the numeric selection.
    """
    while True:
        output_fn("")
        output_fn(f"=== {title} ===")
        for i, (key, desc) in enumerate(options, start=1):
            output_fn(f"  {i}) {key:20s}  {desc}")
        output_fn("")
        output_fn("Choose (number or key): ")
        raw = input_fn("").strip()
        if not raw:
            continue
        if raw.isdigit():
            idx = int(raw) - 1
            if 0 <= idx < len(options):
                return options[idx][0]
        else:
            for key, _ in options:
                if raw == key:
                    return raw
        output_fn(f"Invalid choice {raw!r}, please try again.")


# ---------------------------------------------------------------------------
# Main entry
# ---------------------------------------------------------------------------


def _show_general_info(projectdir: str, project: str, run_nr: str) -> str:
    csv_path = Path(projectdir) / project / run_nr / "0.csv"
    wallclock = calculate_wallclock_time(str(csv_path))
    failed = count_failed_jobs(str(csv_path))
    try:
        total_jobs = sum(1 for _ in csv.DictReader(open(str(csv_path), encoding="utf-8")))
    except FileNotFoundError:
        total_jobs = 0
    return (
        f"Project: {project} (run nr. {run_nr})\n"
        f"Number of jobs (failed and successful): {total_jobs}\n"
        f"Number of failed Jobs: {failed}\n"
        f"Wallclock time: {format_duration(wallclock)}"
    )


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(list(argv) if argv is not None else sys.argv[1:])
    if args.help:
        print(
            "Usage: omniopt_evaluate [OPTIONS]\n"
            "Options:\n"
            "  --projectdir=/path/to/projects/   Path to projects (default: runs)\n"
            "  --nogauge                         Disables gauges\n"
            "  --debug                           Enables debugging\n"
            "  --no_upgrade                      Disables upgrades\n"
            "  --dont_load_modules               Don't load modules (no-op)\n"
            "  --help                            This help"
        )
        return 0

    projectdir = args.projectdir
    if not os.path.isdir(projectdir):
        print(f"Project directory '{projectdir}' does not exist.")
        return 1

    while True:
        projects = find_projects(projectdir)
        if not projects:
            print(f"No projects found in {projectdir}")
            return 1

        menu_options = [(p, "") for p in projects] + [
            ("c", "Change the project dir"),
            ("S", "Start http-server here"),
            ("v", "Show/Change Variables"),
            ("q", "quit"),
        ]
        chosen = textual_menu(f"Available projects under {projectdir}", menu_options)

        if chosen == "q":
            return 0
        if chosen == "c":
            new_dir = textual_input(
                "Projectdir",
                default=projectdir,
            )
            if os.path.isdir(new_dir):
                projectdir = new_dir
            else:
                print(f"'{new_dir}' is not an existing directory.")
            continue
        if chosen == "v":
            print("(variable editor not implemented in textual version yet)")
            continue
        if chosen == "S":
            print("(web-server spin-up not implemented in textual version yet)")
            continue

        # A project was chosen -> show its sub-menu.
        project = chosen
        run_numbers = find_run_numbers(os.path.join(projectdir, project))
        if not run_numbers:
            print(f"No runs found for project {project}.")
            continue
        if len(run_numbers) == 1:
            run_nr = run_numbers[0]
        else:
            opts = [(n, "") for n in run_numbers] + [("b", "back"), ("q", "quit")]
            sub = textual_menu(f"Runs for {project}", opts)
            if sub == "q":
                return 0
            if sub == "b":
                continue
            run_nr = sub

        # Show the info for the chosen run and go back to project list.
        print(_show_general_info(projectdir, project, run_nr))


if __name__ == "__main__":
    sys.exit(main())
