"""The `fluksio sync`, `run`, `runs` and `login` commands.

Kept beside the SDK rather than in `fluksio.cli`: these are the client half of
the tool, and none of them needs the engine to be importable.
"""

from __future__ import annotations

import argparse
import getpass
import importlib
import json
import pkgutil
import sys
from pathlib import Path
from typing import Any

from fluksio.sdk import FLOWS, Flow, SyncError
from fluksio.sdk.client import (
    GLOBAL_DATA_DIR,
    ApiError,
    Client,
    data_dir,
    login,
    origin_of,
    repo_root,
    sync,
)

__all__ = ["add_parsers", "discover"]


def _say(message: str = "") -> None:
    print(message)


def _fail(message: str) -> int:
    print(f"fluksio: {message}", file=sys.stderr)
    return 1


# ---------------------------------------------------------------------------
# Discovery
# ---------------------------------------------------------------------------


def _package_of(directory: Path) -> tuple[str, str]:
    """The path root and dotted name of a package directory."""
    parts = [directory.name]
    parent = directory.parent
    while (parent / "__init__.py").exists():
        parts.append(parent.name)
        parent = parent.parent
    return str(parent), ".".join(reversed(parts))


def _module_of(path: Path) -> tuple[str, str]:
    """The path root and dotted name of a module file."""
    parts = [path.stem]
    directory = path.parent
    while (directory / "__init__.py").exists():
        parts.append(directory.name)
        directory = directory.parent
    return str(directory), ".".join(reversed(parts))


def _import(root: str, dotted: str) -> None:
    if root not in sys.path:
        sys.path.insert(0, root)
    importlib.import_module(dotted)


def discover(targets: list[str]) -> list[Flow]:
    """Import what was named and hand back the flows it declared.

    Imported by dotted name with its root on the path, never from a file
    location: the generated node bodies import the same way, and a module
    loaded under a different name would generate an import that does not
    resolve.
    """
    for target in targets:
        path = Path(target)
        if not path.exists():
            _import(str(Path.cwd()), target)
            continue
        path = path.resolve()
        if path.is_file():
            _import(*_module_of(path))
            continue
        if (path / "__init__.py").exists():
            root, dotted = _package_of(path)
            _import(root, dotted)
            package = sys.modules[dotted]
            for info in pkgutil.walk_packages(package.__path__, f"{dotted}."):
                importlib.import_module(info.name)
            continue
        # A plain directory — a repository root, usually. Its own modules,
        # and the packages inside it: `myresearch/` beside a `README` is the
        # ordinary shape, and naming it explicitly should not be the price of
        # keeping your code in a package.
        for module in sorted(path.glob("*.py")):
            _import(*_module_of(module))
        for child in sorted(path.iterdir()):
            if child.name.startswith(".") or not (child / "__init__.py").exists():
                continue
            root, dotted = _package_of(child)
            _import(root, dotted)
            for info in pkgutil.walk_packages(
                sys.modules[dotted].__path__, f"{dotted}."
            ):
                importlib.import_module(info.name)
    return list(FLOWS.values())


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------


def cmd_login(args: argparse.Namespace) -> int:
    """Only needed for an engine somewhere else — `serve` signs you in here."""
    email = args.email or input("Email: ")
    password = args.password or getpass.getpass("Password: ")
    directory = GLOBAL_DATA_DIR.expanduser() if args.shared else data_dir()
    try:
        path = login(args.url, email, password, directory=directory)
    except ApiError as exc:
        return _fail(f"could not log in: {exc.detail}")
    _say(f"Logged in to {args.url}; the token is in {path}.")
    return 0


def cmd_sync(args: argparse.Namespace) -> int:
    targets = args.targets or ["."]
    try:
        flows = discover(targets)
    except (ImportError, SyncError) as exc:
        return _fail(str(exc))
    if not flows:
        return _fail(
            f"no flows declared in {', '.join(targets)} — a flow is a `Flow(...)` "
            "at module level. Name the package if it is somewhere else: "
            "`fluksio sync src/myresearch`."
        )

    repo = repo_root(targets[0])
    origin = origin_of(repo)
    if origin["dirty"]:
        _say(f"warning: {repo} has uncommitted changes, so the stamp says -dirty")
    if not origin["commit"]:
        _say(f"warning: {repo} is not a git repository, so runs cannot name a commit")

    if args.dry_run:
        for target in flows:
            _say(f"=== flow {target.name}")
            _say(json.dumps(target.document(origin), indent=2))
            for node_id, code in target.shims().items():
                _say(f"=== {target.name}.{node_id}")
                _say(code)
        return 0

    try:
        client = Client(url=args.url, token=args.token)
        reports = sync(
            flows,
            client,
            origin=origin,
            publish=not args.no_publish,
            force=args.force,
        )
    except (SyncError, ApiError) as exc:
        return _fail(str(exc))

    for report in reports:
        if report.unchanged:
            _say(f"  {report.flow}: unchanged")
            continue
        what = "created" if report.created else "updated"
        detail = ", ".join(report.changed)
        state = "published" if report.published else "draft"
        _say(f"  {report.flow}: {what} ({detail}) — {state}")
    stamp = origin["commit"][:7] + ("-dirty" if origin["dirty"] else "")
    _say(f"Stamped with {stamp or 'no commit'} from {repo}.")
    return 0


def _coerce(value: str, dtype: str) -> Any:
    if dtype == "int":
        return int(value)
    if dtype == "float":
        return float(value)
    if dtype == "bool":
        return value.lower() in ("true", "1", "yes", "on")
    if dtype == "str":
        return value
    return json.loads(value)


def _params(definition: dict[str, Any], rest: list[str]) -> dict[str, Any]:
    """Turn `--lr 0.05` into a typed parameter, using the flow's own inputs."""
    types = {
        str(entry["spec"]["name"]): str(entry["spec"].get("dtype", "float"))
        for entry in definition.get("inputs") or []
    }
    params: dict[str, Any] = {}
    pending: str | None = None
    for token in rest:
        if token.startswith("--"):
            if pending is not None:
                # A flag with no value is a flag: `--resume` means true.
                params[pending] = True
            name, sep, value = token[2:].partition("=")
            # Only the name is spelled with dashes; a value may hold one, and
            # `--lr=1e-4` is the case that says so.
            pending = name.replace("-", "_")
            if sep:
                params[pending] = _coerce(value, types.get(pending, "json"))
                pending = None
            continue
        if pending is None:
            raise SyncError(f"unexpected argument '{token}'")
        params[pending] = _coerce(token, types.get(pending, "json"))
        pending = None
    if pending is not None:
        params[pending] = True
    unknown = sorted(set(params) - set(types))
    if unknown:
        raise SyncError(
            f"'{unknown[0]}' is not an input of this flow (it takes "
            f"{', '.join(sorted(types)) or 'none'})"
        )
    return params


def _sync_first(client: Client) -> None:
    """Upload what the working directory declares, before running it.

    The reason a run exists is usually the edit that came before it, and
    remembering to sync is remembering to do the thing the computer could have
    done. So `run` syncs by default — including the worker refresh, which is
    what makes an edit to your own package take effect at all.

    A directory that declares nothing is not an error: a flow drawn on the
    canvas is run the same way, and has nothing to upload.
    """
    try:
        flows = discover(["."])
    except (ImportError, SyncError) as exc:
        # Do not fail a run for a module the run may not even need.
        _say(f"warning: nothing synced — {exc}")
        return
    if not flows:
        return
    repo = repo_root(".")
    reports = sync(flows, client, origin=origin_of(repo))
    changed = [r for r in reports if not r.unchanged]
    if changed:
        _say(f"synced {', '.join(r.flow for r in changed)}")


def cmd_run(args: argparse.Namespace, rest: list[str]) -> int:
    try:
        client = Client(url=args.url, token=args.token)
        if not args.no_sync:
            _sync_first(client)
        stored = client.get_flow(args.flow)
        if stored is None:
            return _fail(f"no flow '{args.flow}' on that engine")
        params = _params(stored.get("definition") or {}, rest)
        handle = client.submit(args.flow, params, seed=args.seed)
    except (SyncError, ApiError) as exc:
        return _fail(str(exc))
    _say(f"{handle.id}  queued  {json.dumps(params)}")
    if not args.wait:
        return 0
    handle.wait(timeout=args.timeout)
    _say(f"{handle.id}  {handle.status}  {json.dumps(handle.result)}")
    return 0 if handle.status == "ok" else 1


def cmd_runs(args: argparse.Namespace) -> int:
    try:
        rows = Client(url=args.url, token=args.token).runs(
            flow=args.flow, limit=args.limit
        )
    except (SyncError, ApiError) as exc:
        return _fail(str(exc))
    for row in rows:
        commit = (row.get("origin_commit") or "")[:7]
        _say(
            f"{row['id']}  {row['status']:<9} {row['flow']:<16} "
            f"{row['duration_ms'] / 1000:7.1f}s  {commit:<8} {json.dumps(row['params'])}"
        )
    return 0


# ---------------------------------------------------------------------------
# Wiring
# ---------------------------------------------------------------------------


def add_parsers(subparsers: Any) -> None:
    """Register the client commands on `fluksio`'s parser."""

    def with_engine(sub: argparse.ArgumentParser) -> None:
        sub.add_argument(
            "--url", default="", help="the engine (default: the last login)"
        )
        sub.add_argument("--token", default="", help="override the stored token")

    parser = subparsers.add_parser(
        "login", help="store a token for an engine elsewhere"
    )
    parser.add_argument("--url", default="http://localhost:8000")
    parser.add_argument("--email", default="")
    parser.add_argument("--password", default="")
    parser.add_argument(
        "--global",
        dest="shared",
        action="store_true",
        help="store it for the machine rather than this project",
    )
    parser.set_defaults(func=cmd_login)

    parser = subparsers.add_parser(
        "sync", help="upload the flows declared in your own code"
    )
    parser.add_argument(
        "targets",
        nargs="*",
        help="modules, packages or directories to import (default: .)",
    )
    parser.add_argument(
        "--dry-run", action="store_true", help="print what would be uploaded"
    )
    parser.add_argument(
        "--no-publish", action="store_true", help="leave the changes as a draft"
    )
    parser.add_argument(
        "--force", action="store_true", help="overwrite work done on the canvas"
    )
    with_engine(parser)
    parser.set_defaults(func=cmd_sync)

    parser = subparsers.add_parser(
        "run", help="sync this directory, then start a run of one of its flows"
    )
    parser.add_argument("flow")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--wait", action="store_true", help="block until it finishes")
    parser.add_argument("--timeout", type=float, default=0.0)
    parser.add_argument(
        "--no-sync",
        action="store_true",
        help="run what is already on the engine, without uploading first",
    )
    with_engine(parser)
    parser.set_defaults(func=cmd_run)

    parser = subparsers.add_parser("runs", help="the runs an engine has recorded")
    parser.add_argument("--flow", default="")
    parser.add_argument("--limit", type=int, default=20)
    with_engine(parser)
    parser.set_defaults(func=cmd_runs)