"""CLI for inspecting and editing mt_dao CURRENT_SERVER_STATE files.

Usage examples (Linux/bash):

    uv run \
        python -m mt_dao.tools.server_state_tool show /path/to/current_server_state.bin

    uv run \
        python -m mt_dao.tools.server_state_tool validate \
        /path/to/current_server_state.bin

    uv run \
        python -m mt_dao.tools.server_state_tool set-round \
        /path/to/current_server_state.bin 42 --inplace

    uv run \
        python -m mt_dao.tools.server_state_tool bump-round \
        /path/to/current_server_state.bin 1 --inplace

    uv run \
        python -m mt_dao.tools.server_state_tool set-time-offset \
        /path/to/current_server_state.bin 123.4 --inplace

    uv run \
        python -m mt_dao.tools.server_state_tool set-server-steps \
        /path/to/current_server_state.bin 10000 --inplace

    uv run \
        python -m mt_dao.tools.server_state_tool set-client-steps \
        /path/to/current_server_state.bin 7 2048 --inplace

    uv run \
        python -m mt_dao.tools.server_state_tool reset-clients \
        /path/to/current_server_state.bin --steps 0 --inplace
"""

from __future__ import annotations

import argparse
import ast
import contextlib
import json
import logging
import pickle  # noqa: S403
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast

import boto3
import fsspec
from botocore.exceptions import ClientError
from fsspec.core import url_to_fs

try:  # Use Flower's logging if available
    from flwr.common import log as _flower_log

    _HAS_FLOWER = True
except ImportError:  # pragma: no cover - fallback in lean environments
    _flower_log = None  # type: ignore[assignment]
    _HAS_FLOWER = False

if TYPE_CHECKING:  # typing-only imports
    from collections.abc import Callable

    from botocore.client import BaseClient

from mt_dao.constants import (
    CLIENT_STATES,
    CURRENT_SERVER_STATE_FILENAME,
    HISTORY,
    LOCAL_STEPS_CUMULATIVE,
    SERVER_ROUND,
    SERVER_STEPS_CUMULATIVE,
    TIME_OFFSET,
)
from mt_dao.server.server_state import check_server_state_lightweight_dict


def mt_dao_log(level: int, msg: str, *args: object) -> None:
    """Lightweight logger that delegates to Flower or stdlib logging."""
    if _HAS_FLOWER and _flower_log is not None:
        _flower_log(level, msg, *args)
    else:
        logging.getLogger("mt_dao").log(level, msg, *args)


# -----------------------------
# I/O helpers
# -----------------------------


def _backup_if_exists(url: str) -> None:
    """Create a .bak backup of ``url`` if it exists.

    Supports both local paths and s3:// URLs using fsspec or boto3 fallback.
    Silently does nothing if the target doesn't exist.
    """
    bak_url = url + ".bak"
    try:
        fs, src = url_to_fs(url)
        if not fs.exists(src):  # nothing to back up
            return
        if url.startswith("s3://") and "://" in bak_url:
            # Prefer boto3 fallback to avoid s3fs dependency
            bucket, key = _parse_s3_url(url)
            data = _s3_get_object_bytes(bucket, key)
            b_bak, k_bak = _parse_s3_url(bak_url)
            _s3_put_object_bytes(b_bak, k_bak, data)
        else:
            with fs.open(src, "rb") as rf, fsspec.open(bak_url, "wb") as wf:
                cast("Any", wf).write(rf.read())
    except (ImportError, FileNotFoundError, ValueError):
        # If fsspec backend is missing but it's S3, fall back to boto3
        if url.startswith("s3://"):
            with contextlib.suppress(FileNotFoundError):
                bucket, key = _parse_s3_url(url)
                data = _s3_get_object_bytes(bucket, key)
                b_bak, k_bak = _parse_s3_url(bak_url)
                _s3_put_object_bytes(b_bak, k_bak, data)
        # else: nothing to back up or unsupported URL


def load_state(path: str | Path) -> dict[str, Any]:
    """Load a CURRENT_SERVER_STATE pickle into a dict.

    Returns
    -------
    dict[str, Any]
        The loaded lightweight server state dictionary.


    Raises
    ------
    FileNotFoundError
        If the provided path does not exist.
    TypeError
        If the unpickled object is not a dictionary.
    ImportError
        If required filesystem drivers are missing for non-S3 remote URLs.
    ValueError
        If parsing the remote URL fails and cannot be handled.

    """
    # Use fsspec so remote URLs (e.g., s3://) work out of the box.
    path_str = str(path)
    try:
        fs, url = url_to_fs(path_str)
        if not fs.exists(url):  # pragma: no cover - safety check
            msg = f"State file not found: {path_str}"
            raise FileNotFoundError(msg)
        with fsspec.open(path_str, "rb") as f:
            state = pickle.load(cast("Any", f))  # noqa: S301
    except (ImportError, ValueError):  # s3fs may be missing
        if path_str.startswith("s3://"):
            bucket, key = _parse_s3_url(path_str)
            _ensure_s3_exists(bucket, key)
            state = pickle.loads(_s3_get_object_bytes(bucket, key))  # noqa: S301
        else:
            raise
    if not isinstance(state, dict):  # pragma: no cover - strictness
        msg = "Expected pickled lightweight dict for server state"
        raise TypeError(msg)
    return state


def save_state(
    state: dict[str, Any], *, path: str | Path, inplace: bool, backup: bool,
) -> str:
    """Save the state dict to disk, optionally in-place with backup.

    Returns
    -------
    str
        The output path (URL) the state was written to.

    Raises
    ------
    ImportError
        If required filesystem drivers are missing for non-S3 remote URLs.

    """
    in_url = str(path)
    # Compute output URL
    if inplace:
        out_url = in_url
    elif in_url.endswith(".bin"):
        # Place edited file next to input, replacing .bin with .edited.bin if present
        out_url = in_url[: -len(".bin")] + ".edited.bin"
    else:
        out_url = in_url + ".edited.bin"
    # Write backups if requested:
    # - In-place: back up the source file
    # - Non-inplace: back up the destination if it already exists
    if backup:
        target_for_backup = in_url if inplace else out_url
        _backup_if_exists(target_for_backup)
    # Ensure local parent exists for file:// or plain local paths
    if "://" not in out_url:
        Path(out_url).parent.mkdir(parents=True, exist_ok=True)
    try:
        # Atomic local write to reduce corruption risk
        if "://" not in out_url:
            tmp_path = Path(str(out_url) + ".tmp")
            tmp_path.write_bytes(pickle.dumps(state))
            # Atomic-ish replace on local FS
            Path(str(out_url)).parent.mkdir(parents=True, exist_ok=True)
            Path(str(tmp_path)).replace(str(out_url))
        else:
            with fsspec.open(out_url, "wb") as f:
                pickle.dump(state, cast("Any", f))
    except ImportError:
        if out_url.startswith("s3://"):
            bucket, key = _parse_s3_url(out_url)
            _s3_put_object_bytes(bucket, key, pickle.dumps(state))
        else:
            raise
    return out_url


# -----------------------------
# State utilities
# -----------------------------


def ensure_client_states_dict(state: dict[str, Any]) -> dict[str | int, dict[str, Any]]:
    """Return client_states as a dict from the lightweight state.

    The lightweight dict stores CLIENT_STATES as a string representation of a dict.
    This helper parses it into a proper dict if present; otherwise returns an empty
    dict.

    Returns
    -------
    dict[str | int, dict[str, Any]]
        The parsed client states mapping.


    Raises
    ------
    TypeError
        If CLIENT_STATES is not a string or dict.
    ValueError
        If parsing the stringified CLIENT_STATES fails.

    """
    if CLIENT_STATES not in state:
        return {}
    value = state[CLIENT_STATES]
    if isinstance(value, str):
        try:
            return ast.literal_eval(value)
        except (ValueError, SyntaxError) as ex:  # pragma: no cover - defensive
            msg = f"Failed to parse CLIENT_STATES string: {ex}"
            raise ValueError(msg) from ex
    if isinstance(value, dict):  # pragma: no cover - rare, but allow
        return value
    msg = "CLIENT_STATES must be a stringified dict in lightweight state"
    raise TypeError(msg)


def put_client_states_dict(
    state: dict[str, Any], client_states: dict[str | int, dict[str, Any]],
) -> None:
    """Store client_states back into the lightweight dict as a string."""
    state[CLIENT_STATES] = str(client_states)


def summarize_state(state: dict[str, Any]) -> dict[str, Any]:
    """Produce a JSON-serializable summary of the lightweight state.

    Returns
    -------
    dict[str, Any]
        A compact, JSON-serializable summary including key fields and a client
    preview.

    """
    clients = ensure_client_states_dict(state)
    sample_clients = list(clients.items())[:5]
    return {
        "file_schema": CURRENT_SERVER_STATE_FILENAME,
        "keys": sorted(state.keys()),
        "server_round": state.get(SERVER_ROUND),
        "time_offset": state.get(TIME_OFFSET),
        "server_steps_cumulative": state.get(SERVER_STEPS_CUMULATIVE),
        "history_type": (
            type(state.get(HISTORY)).__name__ if HISTORY in state else None
        ),
        "n_clients": len(clients),
        "client_sample_preview": sample_clients,
    }


def set_server_round(state: dict[str, Any], new_round: int) -> None:
    """Set the server round value."""
    state[SERVER_ROUND] = int(new_round)


def bump_server_round(state: dict[str, Any], delta: int) -> None:
    """Increment the server round by ``delta`` (can be negative)."""
    state[SERVER_ROUND] = int(state.get(SERVER_ROUND, 0)) + int(delta)


def set_time_offset(state: dict[str, Any], new_offset: float) -> None:
    """Set the time_offset value (seconds)."""
    state[TIME_OFFSET] = float(new_offset)


def set_server_steps(state: dict[str, Any], new_steps: int) -> None:
    """Set the cumulative server steps count."""
    state[SERVER_STEPS_CUMULATIVE] = int(new_steps)


def set_client_steps(state: dict[str, Any], client_id: str | int, steps: int) -> None:
    """Set LOCAL_STEPS_CUMULATIVE for a client, creating if missing."""
    clients = ensure_client_states_dict(state)
    key: str | int = client_id
    entry = clients.get(key, {})
    entry[LOCAL_STEPS_CUMULATIVE] = int(steps)
    clients[key] = entry
    put_client_states_dict(state, clients)


def reset_clients(state: dict[str, Any], steps: int) -> None:
    """Reset all clients' LOCAL_STEPS_CUMULATIVE to ``steps``."""
    clients = ensure_client_states_dict(state)
    for k in list(clients.keys()):
        entry = clients.get(k, {})
        entry[LOCAL_STEPS_CUMULATIVE] = int(steps)
        clients[k] = entry
    put_client_states_dict(state, clients)


def export_json(
    state: dict[str, Any], out_path: str | Path, *, include_history: bool = False,
) -> Path:
    """Export a JSON view of the state for inspection.

    HISTORY may not be JSON-serializable; by default we omit it. When requested,
    it is stringified via repr().

    Returns
    -------
    Path
        Local path object for the written JSON file. For remote URLs, this will
        be a Path constructed from the URL string.

    Raises
    ------
    ImportError
        If required filesystem drivers are missing for non-S3 remote URLs.

    """
    serializable: dict[str, Any] = {k: v for k, v in state.items() if k != HISTORY}
    if include_history and HISTORY in state:
        serializable[HISTORY] = repr(state[HISTORY])
    # Ensure client states become objects (parse from string)
    if CLIENT_STATES in serializable and isinstance(serializable[CLIENT_STATES], str):
        with contextlib.suppress(ValueError, SyntaxError):
            serializable[CLIENT_STATES] = ast.literal_eval(
                serializable[CLIENT_STATES],
            )
    out_url = str(out_path)
    if "://" not in out_url:
        Path(out_url).parent.mkdir(parents=True, exist_ok=True)
    try:
        with fsspec.open(out_url, "w", encoding="utf-8") as fh:
            cast("Any", fh).write(
                json.dumps(serializable, indent=2, sort_keys=True),
            )
    except ImportError:
        if out_url.startswith("s3://"):
            bucket, key = _parse_s3_url(out_url)
            _s3_put_object_bytes(
                bucket,
                key,
                json.dumps(serializable, indent=2, sort_keys=True).encode("utf-8"),
            )
        else:
            raise
    return Path(out_url)


# -----------------------------
# CLI wiring
# -----------------------------


def _add_common_write_args(p: argparse.ArgumentParser) -> None:
    p.add_argument(
        "state_path",
        type=str,
        help="Path to CURRENT_SERVER_STATE file (.bin)",
    )
    p.add_argument("--inplace", action="store_true", help="Modify the file in place")
    p.add_argument(
        "--no-backup",
        action="store_true",
        help="When editing in place, skip writing a .bak backup",
    )


def build_arg_parser() -> argparse.ArgumentParser:
    """Build the CLI argument parser with subcommands.

    Returns
    -------
    argparse.ArgumentParser
        Configured parser instance.

    """
    ap = argparse.ArgumentParser(description=__doc__)
    sub = ap.add_subparsers(dest="cmd", required=True)

    # show
    p_show = sub.add_parser("show", help="Pretty-print a summary of the state")
    p_show.add_argument("state_path", type=str)
    p_show.add_argument(
        "--json",
        dest="json_out",
        type=str,
        help="Optional path to export a JSON summary",
    )

    # validate
    p_val = sub.add_parser(
        "validate",
        help="Validate that the state is lightweight and consistent",
    )
    p_val.add_argument("state_path", type=str)
    p_val.add_argument(
        "--strict",
        action="store_true",
        help="Raise error on invalid state",
    )

    # export-json
    p_exp = sub.add_parser(
        "export-json",
        help="Export the full state to JSON for inspection",
    )
    p_exp.add_argument("state_path", type=str)
    p_exp.add_argument("out", type=str)
    p_exp.add_argument("--include-history", action="store_true")

    # set/bump round, time, steps
    p_sr = sub.add_parser("set-round", help="Set server round to a specific value")
    _add_common_write_args(p_sr)
    p_sr.add_argument("round", type=int)

    p_br = sub.add_parser(
        "bump-round",
        help="Increment server round by delta (can be negative)",
    )
    _add_common_write_args(p_br)
    p_br.add_argument("delta", type=int)

    p_to = sub.add_parser("set-time-offset", help="Set the server time_offset value")
    _add_common_write_args(p_to)
    p_to.add_argument("time_offset", type=float)

    p_ss = sub.add_parser(
        "set-server-steps",
        help="Set SERVER_STEPS_CUMULATIVE",
    )
    _add_common_write_args(p_ss)
    p_ss.add_argument("steps", type=int)

    # client ops
    p_cs = sub.add_parser(
        "set-client-steps",
        help="Set LOCAL_STEPS_CUMULATIVE for a client",
    )
    _add_common_write_args(p_cs)
    p_cs.add_argument("client_id", type=str, help="Client id (string or integer)")
    p_cs.add_argument("steps", type=int)

    p_rc = sub.add_parser(
        "reset-clients",
        help="Reset all clients' LOCAL_STEPS_CUMULATIVE",
    )
    _add_common_write_args(p_rc)
    p_rc.add_argument("--steps", type=int, default=0)

    # upgrade scaffold
    p_up = sub.add_parser(
        "upgrade",
        help="Upgrade the lightweight state to the latest schema (scaffold)",
    )
    _add_common_write_args(p_up)

    return ap


# -----------------------------
# Upgrade scaffold (extensible migrations)
# -----------------------------


def upgrade_state(state: dict[str, Any]) -> dict[str, Any]:
    """Upgrade a lightweight state in-place to the latest schema.

    This is a scaffold for future migrations. Extend by applying transformations
    when older keys/structures are detected. Currently, it returns the state
    unchanged.

    Returns
    -------
    dict[str, Any]
        The potentially modified state dictionary.

    """
    return state


# -----------------------------
# S3 helpers (boto3 fallback)
# -----------------------------


def _parse_s3_url(url: str) -> tuple[str, str]:
    url_no_scheme = url[len("s3://") :]
    bucket, _, key = url_no_scheme.partition("/")
    return bucket, key


def _s3_client() -> BaseClient:  # pragma: no cover - simple factory
    return cast("BaseClient", boto3.client("s3"))


def _ensure_s3_exists(bucket: str, key: str) -> None:
    try:
        _s3_client().head_object(Bucket=bucket, Key=key)
    except ClientError as ex:  # pragma: no cover - network/env dependent
        msg = f"State file not found: s3://{bucket}/{key} ({ex})"
        raise FileNotFoundError(msg) from ex


def _s3_get_object_bytes(bucket: str, key: str) -> bytes:
    res = _s3_client().get_object(Bucket=bucket, Key=key)
    body = res.get("Body")
    return body.read() if body else b""


def _s3_put_object_bytes(bucket: str, key: str, data: bytes) -> None:
    _s3_client().put_object(Bucket=bucket, Key=key, Body=data)


def _handle_upgrade(args: argparse.Namespace) -> int:
    state = load_state(args.state_path)
    upgraded = upgrade_state(state)
    out = save_state(
        upgraded,
        path=args.state_path,
        inplace=args.inplace,
        backup=not args.no_backup,
    )
    mt_dao_log(logging.INFO, "State upgraded and written: %s", out)
    return 0


def _handle_show(args: argparse.Namespace) -> int:
    state = load_state(args.state_path)
    summary = summarize_state(state)
    mt_dao_log(logging.INFO, json.dumps(summary, indent=2))
    if args.json_out:
        export_json(state, args.json_out)
    return 0


def _handle_validate(args: argparse.Namespace) -> int:
    state = load_state(args.state_path)
    is_ok = check_server_state_lightweight_dict(
        state,
        raise_if_not_lightweight=False,
    )
    mt_dao_log(logging.INFO, "%s", "VALID" if is_ok else "INVALID")
    if args.strict and not is_ok:
        raise SystemExit(2)
    return 0


def _handle_export_json(args: argparse.Namespace) -> int:
    state = load_state(args.state_path)
    out = export_json(state, args.out, include_history=args.include_history)
    mt_dao_log(logging.INFO, "%s", str(out))
    return 0


def _coerce_client_id(cid: str) -> str | int:
    try:
        return int(cid)
    except ValueError:
        return cid


def _handle_edit(
    args: argparse.Namespace,
    editor: Callable[[dict[str, Any]], None],
) -> int:
    state = load_state(args.state_path)
    editor(state)
    out = save_state(
        state,
        path=args.state_path,
        inplace=args.inplace,
        backup=not args.no_backup,
    )
    mt_dao_log(logging.INFO, "State written: %s", str(out))
    return 0


def main(argv: list[str] | None = None) -> int:
    """Program entrypoint for the server state CLI.

    Returns
    -------
    int
        Exit code compatible with shell expectations.

    """
    ap = build_arg_parser()
    args = ap.parse_args(argv)

    dispatch: dict[str, Callable[[argparse.Namespace], int]] = {
        "show": _handle_show,
        "validate": _handle_validate,
        "export-json": _handle_export_json,
        "upgrade": _handle_upgrade,
    }

    exit_code = 0
    if args.cmd in dispatch:
        exit_code = dispatch[args.cmd](args)
    elif args.cmd == "set-round":
        exit_code = _handle_edit(args, lambda s: set_server_round(s, args.round))
    elif args.cmd == "bump-round":
        exit_code = _handle_edit(args, lambda s: bump_server_round(s, args.delta))
    elif args.cmd == "set-time-offset":
        exit_code = _handle_edit(args, lambda s: set_time_offset(s, args.time_offset))
    elif args.cmd == "set-server-steps":
        exit_code = _handle_edit(args, lambda s: set_server_steps(s, args.steps))
    elif args.cmd == "set-client-steps":
        cid = _coerce_client_id(args.client_id)
        exit_code = _handle_edit(args, lambda s: set_client_steps(s, cid, args.steps))
    elif args.cmd == "reset-clients":
        exit_code = _handle_edit(args, lambda s: reset_clients(s, args.steps))
    else:  # pragma: no cover
        ap.error(f"Unknown command: {args.cmd}")
        exit_code = 2

    return exit_code


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())
