#!/usr/bin/env python3
"""omniopt_plot - Dispatcher for OmniOpt run plots.

Python rewrite of the former ``omniopt_plot`` bash script.

Usage:
    omniopt_plot [OPTIONS] [plot_type-or-run-dir]

The script:
  * discovers all plot scripts matching ``.omniopt_plot_*.py``
  * parses each script's ``# DESCRIPTION`` / ``# EXPECTED FILES`` metadata
  * builds a textual menu of plot types whose expected files exist in the
    chosen run directory
  * dispatches to the selected ``.omniopt_plot_<type>.py`` script

CLI flags are 100 % backwards-compatible with the bash version.

Behavioural compatibility:
  * ``--run_dir``, ``--save_to_file``, ``--min``, ``--max``,
    ``--allow_axes``, ``--plot_type``, ``--help``, ``--debug``
  * ``python3 .omniopt_plot_<type>.py ...`` is invoked the same way
  * The textual menu replaces ``whiptail`` (no TUI dependency).
"""

from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Callable, List, Optional, Sequence, Tuple


SCRIPT_DIR = Path(__file__).resolve().parent
MENU_KEY = "menu"

# Default location of plot scripts (sibling files named .omniopt_plot_*.py)
PLOT_SCRIPT_GLOB = ".omniopt_plot_*.py"

ENV_VARS_FOR_PLOT_RUN = (
    "MPLCONFIGDIR",
    "XDG_CACHE_HOME",
    "WHIPTAIL",
)


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


def levenshtein(a: str, b: str) -> int:
    """Standard iterative Levenshtein distance."""
    if a == b:
        return 0
    if not a:
        return len(b)
    if not b:
        return len(a)

    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, start=1):
        cur = [i]
        for j, cb in enumerate(b, start=1):
            cost = 0 if ca == cb else 1
            cur.append(
                min(
                    cur[j - 1] + 1,        # insertion
                    prev[j] + 1,           # deletion
                    prev[j - 1] + cost,    # substitution
                )
            )
        prev = cur
    return prev[-1]


def find_closest_match(user_input: str, candidates: Sequence[str]) -> List[str]:
    """Return the candidate(s) that best match ``user_input``.

    If any candidate contains ``user_input`` as a substring, all those
    candidates are returned (the caller picks).  Otherwise the single
    candidate with the smallest Levenshtein distance is returned.
    """
    substring_matches = [c for c in candidates if user_input in c]
    if len(substring_matches) == 1:
        return substring_matches
    if len(substring_matches) > 1:
        return substring_matches
    if not candidates:
        return []
    best = min(candidates, key=lambda c: levenshtein(user_input, c))
    return [best]


def list_plot_types(script_dir: str) -> List[str]:
    """Return sorted list of plot type names (``scatter``, ``general``, ...)."""
    root = Path(script_dir)
    out: List[str] = []
    for path in root.glob(PLOT_SCRIPT_GLOB):
        name = path.stem
        # .omniopt_plot_<type>.py -> <type>
        prefix = ".omniopt_plot_"
        if name.startswith(prefix):
            out.append(name[len(prefix):])
    return sorted(set(out))


def _read_plot_metadata(plot_type: str, script_dir: str) -> str:
    path = Path(script_dir) / f".omniopt_plot_{plot_type}.py"
    try:
        return path.read_text(encoding="utf-8", errors="replace")
    except FileNotFoundError:
        return ""


def get_plot_description(plot_type: str, script_dir: str) -> str:
    """Return the ``# DESCRIPTION: ...`` line for the given plot type."""
    body = _read_plot_metadata(plot_type, script_dir)
    for line in body.splitlines():
        m = re.match(r"#\s*DESCRIPTION:\s*(.+?)\s*$", line)
        if m:
            return m.group(1)
    return ""


def get_expected_files(plot_type: str, script_dir: str) -> List[str]:
    """Parse ``# EXPECTED FILES: a, b, c`` into a list of substrings."""
    body = _read_plot_metadata(plot_type, script_dir)
    for line in body.splitlines():
        m = re.match(r"#\s*EXPECTED FILES:\s*(.+?)\s*$", line)
        if m:
            return [s.strip() for s in m.group(1).split(",") if s.strip()]
    return []


def plot_type_accepts_min_max(plot_type: str, script_dir: str) -> bool:
    """Heuristic: plot supports ``args.min``/``args.max`` if it references
    ``args.min`` (and the reference is not commented-out / disabled)."""
    body = _read_plot_metadata(plot_type, script_dir)
    if not body:
        return False
    return bool(re.search(r"(?<!useless )\bargs\.min\b", body))


def plot_type_supports_save_to_file(plot_type: str, script_dir: str) -> bool:
    """Plot supports ``--save_to_file`` if it has an add_argument with that name."""
    body = _read_plot_metadata(plot_type, script_dir)
    if not body:
        return False
    return bool(re.search(r"add_argument\b[^\n]*\bsave_to_file\b", body))


def check_plot_prerequisites(plot_type: str, run_dir: str, script_dir: str) -> bool:
    """All expected file-substrings must appear in the run directory."""
    expected = get_expected_files(plot_type, script_dir)
    if not expected:
        return False
    try:
        entries = set(os.listdir(run_dir))
    except (FileNotFoundError, NotADirectoryError):
        return False
    return all(any(eff in entry for entry in entries) for eff in expected)


def validate_plot_type(plot_type: str, script_dir: str) -> str:
    """Return ``""`` on success, error message otherwise."""
    valid = [MENU_KEY] + list_plot_types(script_dir)
    if plot_type in valid:
        return ""
    return f"Invalid plot type {plot_type!r}, valid: {', '.join(valid)}"


def resolve_run_dir(
    run_dir: str, original_pwd: str, in_docker_user_dir: bool
) -> str:
    """Resolve a possibly-relative ``run_dir`` to an absolute path."""
    if not run_dir or os.path.isabs(run_dir):
        return run_dir
    prefix = os.path.join(original_pwd, "docker_user_dir")
    base = prefix if in_docker_user_dir else original_pwd
    return os.path.join(base, run_dir)


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


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="omniopt_plot",
        description="Plot OmniOpt runs.",
        add_help=False,
    )
    p.add_argument("--help", "-h", action="store_true")
    p.add_argument("--debug", action="store_true")
    p.add_argument("--run_dir", default=None)
    p.add_argument("--save_to_file", default="0")
    p.add_argument("--min", default="")
    p.add_argument("--max", default="")
    p.add_argument("--allow_axes", default=None)
    p.add_argument("--plot_type", default=MENU_KEY)
    # Positional argument: either a run_dir or a (fuzzy) plot_type.
    p.add_argument(
        "positional", nargs="?", default=None,
        help="Run directory (if it exists) or a plot type name.",
    )
    return p


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    """Parse argv.  Accepts either ``--run_dir=...`` or a positional arg
    that is either an existing directory (treated as ``run_dir``) or a
    fuzzy match against the available plot types.

    Mirrors the bash version's behaviour so that ``omniopt_plot
    runs/myexp/0`` works just as well as ``omniopt_plot --run_dir
    runs/myexp/0``.
    """
    args = _build_parser().parse_args(list(argv))
    if args.positional is not None and not args.positional.startswith("-"):
        script_dir = str(SCRIPT_DIR)
        plot_types = list_plot_types(script_dir)
        if os.path.isdir(args.positional):
            if args.run_dir is None:
                args.run_dir = args.positional
        else:
            match = find_closest_match(args.positional, plot_types)
            if match:
                args.plot_type = match[0]
    return args


# ---------------------------------------------------------------------------
# Plot dispatch
# ---------------------------------------------------------------------------


def _build_args_string(args: argparse.Namespace) -> str:
    """Reproduce the bash ``args_string`` accumulator."""
    parts: List[str] = []
    if args.run_dir:
        parts.append(f"--run_dir={args.run_dir}")
    if args.save_to_file and args.save_to_file != "0":
        parts.append(f"--save_to_file={args.save_to_file}")
    if args.min:
        parts.append(f"--min={args.min}")
    if args.max:
        parts.append(f"--max={args.max}")
    if args.allow_axes:
        parts.append(f"--allow_axes={args.allow_axes}")
    return " ".join(parts)


def _print_help(script_dir: str, plot_type: str) -> None:
    if plot_type and plot_type != MENU_KEY:
        subprocess.run(
            [sys.executable, f"{script_dir}/.omniopt_plot_{plot_type}.py", "--help"],
            check=False,
        )
    else:
        types = list_plot_types(script_dir)
        print("omniopt_plot: Plot omniopt runs")
        print("Basic usage: omniopt_plot --run_dir=runs/testrun/0 --plot_type=scatter")
        print("Possible options for plot_type:")
        for t in types:
            print(f"  - {t}")
        print("For specific options, use omniopt_plot --plot_type=scatter --help")


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv if argv is not None else sys.argv[1:])
    script_dir = str(SCRIPT_DIR)
    original_pwd = os.environ.get("ORIGINAL_PWD", os.getcwd())

    if args.help:
        _print_help(script_dir, args.plot_type)
        return 0

    err = validate_plot_type(args.plot_type, script_dir)
    if err:
        print(f"Error: {err}", file=sys.stderr)
        return 99

    if args.run_dir:
        args.run_dir = resolve_run_dir(
            args.run_dir, original_pwd, os.path.isdir("docker_user_dir")
        )

    if args.plot_type == MENU_KEY:
        if not args.run_dir:
            print("Error: --run_dir is missing", file=sys.stderr)
            return 1
        if not os.path.isdir(args.run_dir):
            print(f"Error: --run_dir is not a directory: {args.run_dir}", file=sys.stderr)
            return 1
        chosen = _textual_plot_menu(args.run_dir, script_dir)
        if chosen is None:
            return 0
        args.plot_type = chosen

    # Validate file prerequisites
    if args.run_dir:
        if not check_plot_prerequisites(args.plot_type, args.run_dir, script_dir):
            print(
                f"It seems like the run folder {args.run_dir} does not have any plotable data.",
                file=sys.stderr,
            )
            return 3

    # Dispatch
    plot_script = Path(script_dir) / f".omniopt_plot_{args.plot_type}.py"
    args_string = _build_args_string(args)
    cmd = f"{sys.executable} {plot_script} {args_string}".strip()
    env = os.environ.copy()
    env["RUN_VIA_RUNSH"] = "1"
    proc = subprocess.run(cmd, shell=True, env=env)
    return proc.returncode


def _textual_plot_menu(run_dir: str, script_dir: str) -> Optional[str]:
    """Print the textual menu and return the chosen plot type (or None
    when the user quits)."""
    types = list_plot_types(script_dir)
    available: List[Tuple[str, str, bool]] = []
    for t in types:
        if not check_plot_prerequisites(t, run_dir, script_dir):
            continue
        desc = get_plot_description(t, script_dir)
        suffix = ", honors min/max" if plot_type_accepts_min_max(t, script_dir) else ""
        available.append((t, f"{desc}{suffix}", plot_type_supports_save_to_file(t, script_dir)))

    if not available:
        print(
            f"It seems like the run folder {run_dir} does not have any plotable data.",
            file=sys.stderr,
        )
        return None

    print()
    print(f"=== Available plots for {run_dir} ===")
    for i, (key, desc, _) in enumerate(available, start=1):
        print(f"  {i}) {key:20s}  {desc}")
    print(f"  q) quit")
    print()
    while True:
        try:
            raw = input("Choose (number or key): ").strip()
        except EOFError:
            return None
        if raw == "q":
            return None
        if raw.isdigit():
            idx = int(raw) - 1
            if 0 <= idx < len(available):
                return available[idx][0]
        else:
            for key, _, _ in available:
                if raw == key:
                    return key
        print(f"Invalid choice {raw!r}, please try again.")


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