"""Hybrid local/EC2 execution for long-running data-plane commands.

Short, interactive steps run locally (see :mod:`runner`). Long jobs (metadata
upload/delete, synth upload/delete, indexd register) can instead run on the
project's EC2 box via ``--on ec2``. The default mechanism is **AWS SSM Run
Command**, which is disconnect-safe by construction: the SSM agent executes the
job server-side, so the laptop can sleep and you re-attach by ``run_id``.

Authentication split (the crux):
    * the laptop authenticates the SSM call with the *local* env's named profile
      (e.g. ``staging`` -> ``acdc_staging``);
    * the job on the box runs under the ``*_ec2`` env (``aws_profile: null`` ->
      the instance profile).
"""
from __future__ import annotations

import datetime
import shlex
import time
from enum import Enum
from typing import Callable, List, Sequence, Tuple

import typer
from botocore.exceptions import ClientError

from acdc_aws_etl_pipeline.config import (
    ConfigError,
    EnvConfig,
    env_base,
    load_config,
    repo_root,
    resolve_env,
)
from acdc_aws_etl_pipeline.upload.metadata_submitter import create_boto3_session
from acdc_aws_etl_pipeline.cli._internal import registry, runner

#: Where the repo is cloned on the EC2 box (see docs/ec2_setup.md).
REPO_DIR_ON_EC2 = "~/acdc-aws-etl-pipeline"
#: S3 location SSM writes full command output to (sidesteps the 24KB inline cap).
DEFAULT_LOG_BUCKET = "acdc-dataops-metadata"
DEFAULT_LOG_PREFIX = "acdc-runs"
#: CloudWatch Logs group the SSM agent streams dispatched-run output to, so
#: ``acdc jobs logs --follow`` can tail the job in near-real-time (S3 output is
#: only uploaded on completion). Requires CloudWatch Logs write perms on the
#: EC2 instance role (see docs/ec2_setup.md).
DEFAULT_LOG_GROUP = "/acdc/jobs"


class Target(str, Enum):
    """Where a command should execute."""

    local = "local"
    ec2 = "ec2"


class DispatchMech(str, Enum):
    """How an EC2 job is launched."""

    ssm = "ssm"
    ssh = "ssh"


def new_run_id(label: str) -> str:
    """Build a sortable, human-friendly run id: ``20260615T1430-metadata-upload``."""
    ts = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
    return f"{ts}-{label}"


def resolve_dispatch_envs(cfg: dict, env: str) -> Tuple[EnvConfig, EnvConfig]:
    """Return ``(auth_env, remote_env)`` for an EC2 dispatch.

    ``auth_env`` is the non-ec2 profile used to authenticate the SSM/ssh call
    from the laptop; ``remote_env`` is the ``*_ec2`` variant that carries the
    instance id and the ``--env`` the remote job runs under. Accepts either the
    base or the ``_ec2`` form for ``env``.
    """
    base = env_base(env)
    auth_env = resolve_env(cfg, base)
    remote_env = resolve_env(cfg, f"{base}_ec2")
    return auth_env, remote_env


def build_remote_command(run_id: str, remote_argv: Sequence[str]) -> str:
    """Build the shell snippet SSM/ssh runs on the EC2 box."""
    inner = " ".join(str(a) for a in remote_argv)
    return " && ".join(
        [
            "set -euo pipefail",
            f"cd {REPO_DIR_ON_EC2}",
            "git pull --ff-only",
            "mkdir -p logs",
            # Tee to both the on-box archive and stdout (which SSM forwards to
            # CloudWatch Logs for live `acdc jobs logs --follow`). PYTHONUNBUFFERED
            # makes the stream line-by-line; `set -o pipefail` (above) keeps the
            # job's exit status flowing through `tee` so failures still mark Failed.
            f"PYTHONUNBUFFERED=1 poetry run {inner} 2>&1 | tee logs/{run_id}.log",
        ]
    )


def dispatch_ssm(
    auth_env: EnvConfig,
    remote_env: EnvConfig,
    remote_argv: Sequence[str],
    label: str,
    *,
    log_bucket: str = DEFAULT_LOG_BUCKET,
    log_prefix: str = DEFAULT_LOG_PREFIX,
) -> str:
    """Launch ``remote_argv`` on the EC2 box via SSM Run Command; return ``run_id``."""
    if not remote_env.ec2_instance_id:
        typer.secho(
            f"No 'ec2_instance_id' configured for env '{remote_env.name}'. "
            f"Add it to deploy_config.yaml to use --on ec2.",
            fg=typer.colors.RED,
            err=True,
        )
        raise typer.Exit(2)

    session = create_boto3_session(
        aws_profile=auth_env.aws_profile, aws_region=auth_env.region
    )
    ssm = session.client("ssm")
    run_id = new_run_id(label)
    # SSM AWS-RunShellScript executes as root, whose HOME is /root and whose PATH
    # lacks poetry; the repo, poetry, and git credentials all live under ec2-user.
    # Run the job inside that user's login shell so ~, PATH, and git creds resolve.
    remote_user = remote_env.ssh_user or "ec2-user"
    command = "runuser -l {user} -c {script}".format(
        user=remote_user,
        script=shlex.quote(build_remote_command(run_id, remote_argv)),
    )

    resp = ssm.send_command(
        InstanceIds=[remote_env.ec2_instance_id],
        DocumentName="AWS-RunShellScript",
        Parameters={"commands": [command], "executionTimeout": ["172800"]},
        OutputS3BucketName=log_bucket,
        OutputS3KeyPrefix=f"{log_prefix}/{run_id}",
        CloudWatchOutputConfig={
            "CloudWatchLogGroupName": DEFAULT_LOG_GROUP,
            "CloudWatchOutputEnabled": True,
        },
        Comment=f"acdc {run_id}",
    )
    command_id = resp["Command"]["CommandId"]
    s3_log_uri = f"s3://{log_bucket}/{log_prefix}/{run_id}"
    registry.record(
        run_id,
        command_id=command_id,
        instance_id=remote_env.ec2_instance_id,
        env=remote_env.name,
        argv=list(remote_argv),
        mechanism=DispatchMech.ssm.value,
        s3_log_uri=s3_log_uri,
        cw_log_group=DEFAULT_LOG_GROUP,
        started_at=run_id.split("-", 1)[0],
    )
    _print_dispatch_banner(run_id, remote_env)
    return run_id


def dispatch_ssh(
    auth_env: EnvConfig,
    remote_env: EnvConfig,
    remote_argv: Sequence[str],
    label: str,
    *,
    ssh_key: str,
) -> str:
    """Fallback: launch the job over SSH with ``nohup`` (reuses the ec2_setup flow).

    The job survives the SSH session ending; status/logs are polled, not streamed
    over a held connection.
    """
    host = remote_env.ec2_instance_id  # caller may pass a DNS/host via config
    user = remote_env.ssh_user or "ec2-user"
    run_id = new_run_id(label)
    remote = build_remote_command(run_id, remote_argv) + " &"
    ssh_argv = ["ssh", "-i", ssh_key, f"{user}@{host}", remote]
    runner.run(ssh_argv)
    registry.record(
        run_id,
        command_id=None,
        instance_id=host,
        env=remote_env.name,
        argv=list(remote_argv),
        mechanism=DispatchMech.ssh.value,
        s3_log_uri=None,
        started_at=run_id.split("-", 1)[0],
    )
    _print_dispatch_banner(run_id, remote_env)
    return run_id


def run_or_dispatch(
    on: Target,
    cfg: dict,
    env: str,
    script_relpath: str,
    build_args: Callable[[str], Sequence[str]],
    label: str,
    *,
    interpreter: str = "python",
    mech: DispatchMech = DispatchMech.ssm,
    ssh_key: str = None,
) -> None:
    """Run a service script locally, or dispatch it to EC2.

    ``build_args(env_name)`` returns the service-script arguments for the given
    ``--env`` value. Locally we run ``<interpreter> <script> <args>`` with the
    chosen env; for EC2 we dispatch the same with the ``*_ec2`` env.

    ``interpreter`` is ``"python"`` (run with the venv interpreter) or ``"bash"``.

    Note: the remote argv targets the existing ``services/*`` script (not the
    ``acdc`` CLI) so it works before the EC2 box is re-bootstrapped with this
    package version.
    """
    try:
        if on == Target.local:
            local_env = resolve_env(cfg, env)
            args = build_args(local_env.name)
            if interpreter == "bash":
                runner.run(runner.bash_script(script_relpath, *args))
            else:
                runner.run(runner.python_script(script_relpath, *args))
            return
        auth_env, remote_env = resolve_dispatch_envs(cfg, env)
    except ConfigError as exc:
        typer.secho(str(exc), fg=typer.colors.RED, err=True)
        raise typer.Exit(1)

    remote_argv = [interpreter, script_relpath, *build_args(remote_env.name)]
    if mech == DispatchMech.ssh:
        if not ssh_key and not remote_env.ssh_key:
            typer.secho(
                "--dispatch ssh requires --key or ssh_key in config.",
                fg=typer.colors.RED,
                err=True,
            )
            raise typer.Exit(2)
        dispatch_ssh(
            auth_env, remote_env, remote_argv, label,
            ssh_key=ssh_key or remote_env.ssh_key,
        )
    else:
        dispatch_ssm(auth_env, remote_env, remote_argv, label)


# --------------------------------------------------------------------------- #
# Watching dispatched runs (acdc jobs)                                         #
# --------------------------------------------------------------------------- #
def _auth_session_for_run(rec: dict):
    cfg = load_config()
    auth_env = resolve_env(cfg, env_base(rec["env"]))
    return create_boto3_session(
        aws_profile=auth_env.aws_profile, aws_region=auth_env.region
    )


def status(run_id: str) -> dict:
    """Return the SSM invocation status for a dispatched run."""
    rec = registry.get(run_id)
    if not rec:
        typer.secho(f"Unknown run id: {run_id}", fg=typer.colors.RED, err=True)
        raise typer.Exit(1)
    if rec.get("mechanism") != DispatchMech.ssm.value or not rec.get("command_id"):
        typer.secho(
            f"Run {run_id} was not dispatched via SSM; check it on the box.",
            fg=typer.colors.YELLOW,
        )
        return rec
    ssm = _auth_session_for_run(rec).client("ssm")
    return ssm.get_command_invocation(
        CommandId=rec["command_id"], InstanceId=rec["instance_id"]
    )


def status_label(rec: dict) -> str:
    """Best-effort current status for the all-runs listing (never raises).

    Unlike :func:`status`, this is keyed off the registry record and stays quiet so
    it can be called in a loop. SSH-dispatched runs have no SSM invocation to query.
    """
    if rec.get("mechanism") != DispatchMech.ssm.value or not rec.get("command_id"):
        return "n/a (ssh)"
    ssm = _auth_session_for_run(rec).client("ssm")
    try:
        inv = ssm.get_command_invocation(
            CommandId=rec["command_id"], InstanceId=rec["instance_id"]
        )
        return inv.get("Status", "unknown")
    except ssm.exceptions.InvocationDoesNotExist:
        return "pending"  # sent but not yet registered, or aged out of SSM history
    except ClientError:
        return "unknown"


def stop(run_id: str) -> None:
    """Cancel a running SSM-dispatched job.

    Asks the SSM agent to terminate the in-flight command on the box. Cancellation
    isn't instantaneous and isn't guaranteed by AWS, but in practice it stops the
    process; ``acdc jobs status <run_id>`` shows ``Cancelled`` once it takes effect.
    """
    rec = registry.get(run_id)
    if not rec:
        typer.secho(f"Unknown run id: {run_id}", fg=typer.colors.RED, err=True)
        raise typer.Exit(1)
    if rec.get("mechanism") != DispatchMech.ssm.value or not rec.get("command_id"):
        typer.secho(
            f"Run {run_id} was not dispatched via SSM and can't be stopped remotely; "
            f"ssh to the box and kill it.",
            fg=typer.colors.YELLOW,
            err=True,
        )
        raise typer.Exit(1)
    current = status_label(rec)
    if current in _TERMINAL_STATUSES:
        typer.secho(f"Run {run_id} already finished: {current}", fg=typer.colors.YELLOW)
        return
    ssm = _auth_session_for_run(rec).client("ssm")
    try:
        ssm.cancel_command(
            CommandId=rec["command_id"], InstanceIds=[rec["instance_id"]]
        )
    except ClientError as exc:
        typer.secho(f"Could not stop {run_id}: {exc}", fg=typer.colors.RED, err=True)
        raise typer.Exit(1)
    typer.secho(
        f"Stop requested for {run_id} ({rec['instance_id']}).", fg=typer.colors.GREEN
    )


def _read_s3_logs(rec: dict) -> str:
    """Best-effort fetch of the full stdout/stderr SSM wrote to S3."""
    uri = rec.get("s3_log_uri") or ""
    if not uri.startswith("s3://"):
        return ""
    bucket, _, prefix = uri[len("s3://"):].partition("/")
    s3 = _auth_session_for_run(rec).client("s3")
    listing = s3.list_objects_v2(Bucket=bucket, Prefix=prefix)
    chunks: List[str] = []
    for obj in listing.get("Contents", []):
        key = obj["Key"]
        if key.endswith(("stdout", "stderr")):
            body = s3.get_object(Bucket=bucket, Key=key)["Body"].read()
            chunks.append(body.decode("utf-8", errors="replace"))
    return "\n".join(chunks)


#: SSM invocation states that mean the run is over.
_TERMINAL_STATUSES = ("Success", "Failed", "Cancelled", "TimedOut")


def _read_cw_logs(rec: dict, start_time: int = 0) -> List[dict]:
    """Fetch CloudWatch Logs events for a run's SSM stdout/stderr streams.

    SSM names the streams ``<command-id>/<instance-id>/<plugin>/{stdout,stderr}``,
    so the ``<command-id>/<instance-id>`` prefix captures both. ``start_time`` (epoch
    millis) lets the follow loop pull only newer events; callers de-dup by ``eventId``.
    Returns ``[]`` if the group/stream doesn't exist yet (run just started, or the
    instance role lacks CloudWatch Logs write permission).
    """
    group = rec.get("cw_log_group")
    command_id = rec.get("command_id")
    instance_id = rec.get("instance_id")
    if not (group and command_id and instance_id):
        return []
    client = _auth_session_for_run(rec).client("logs")
    kwargs = {
        "logGroupName": group,
        "logStreamNamePrefix": f"{command_id}/{instance_id}",
        "startTime": start_time,
    }
    events: List[dict] = []
    try:
        while True:
            resp = client.filter_log_events(**kwargs)
            events.extend(resp.get("events", []))
            token = resp.get("nextToken")
            if not token:
                break
            kwargs["nextToken"] = token
    except client.exceptions.ResourceNotFoundException:
        return []
    return events


def _stream_cw_logs(run_id: str, rec: dict, *, follow: bool, poll_seconds: int) -> None:
    """Print a run's CloudWatch output, optionally following until it finishes."""
    seen_ids: set = set()
    last_ts = 0

    def drain() -> None:
        nonlocal last_ts
        for ev in sorted(_read_cw_logs(rec, last_ts), key=lambda e: e["timestamp"]):
            if ev["eventId"] in seen_ids:
                continue
            seen_ids.add(ev["eventId"])
            last_ts = max(last_ts, ev["timestamp"])
            typer.echo(ev["message"])

    while True:
        drain()
        if not follow:
            break
        inv = status(run_id)
        if isinstance(inv, dict) and inv.get("Status") in _TERMINAL_STATUSES:
            drain()  # catch events ingested between the last drain and completion
            typer.secho(f"\n[run {run_id} finished: {inv['Status']}]",
                        fg=typer.colors.BRIGHT_BLACK)
            break
        time.sleep(poll_seconds)


def logs(run_id: str, follow: bool = False, poll_seconds: int = 5) -> None:
    """Print (and optionally follow) the logs for a dispatched run."""
    rec = registry.get(run_id)
    if not rec:
        typer.secho(f"Unknown run id: {run_id}", fg=typer.colors.RED, err=True)
        raise typer.Exit(1)

    # Runs dispatched with CloudWatch output stream live; older runs only wrote to
    # S3 (uploaded on completion), so fall back to that.
    if rec.get("cw_log_group"):
        _stream_cw_logs(run_id, rec, follow=follow, poll_seconds=poll_seconds)
        return

    seen = 0
    while True:
        text = _read_s3_logs(rec)
        if len(text) > seen:
            typer.echo(text[seen:], nl=False)
            seen = len(text)
        if not follow:
            break
        inv = status(run_id)
        if isinstance(inv, dict) and inv.get("Status") in _TERMINAL_STATUSES:
            typer.secho(f"\n[run {run_id} finished: {inv['Status']}]",
                        fg=typer.colors.BRIGHT_BLACK)
            break
        time.sleep(poll_seconds)


def _print_dispatch_banner(run_id: str, remote_env: EnvConfig) -> None:
    typer.secho(
        f"Dispatched to EC2 ({remote_env.ec2_instance_id} / {remote_env.name})",
        fg=typer.colors.GREEN,
    )
    typer.echo(f"  run id : {run_id}")
    typer.echo(f"  logs   : acdc jobs logs {run_id} --follow")
    typer.echo(f"  status : acdc jobs status {run_id}")


# Keep a reference to repo_root so callers importing dispatch have it handy.
__all__ = [
    "Target",
    "DispatchMech",
    "run_or_dispatch",
    "dispatch_ssm",
    "dispatch_ssh",
    "resolve_dispatch_envs",
    "build_remote_command",
    "new_run_id",
    "status",
    "status_label",
    "stop",
    "logs",
    "repo_root",
]
