"""abreka exp remote — inspect and recover remote-only run execution."""

from __future__ import annotations

import json
import time
from datetime import datetime, timezone
from pathlib import Path

import click
from rich.console import Console

from abreka.config import find_abreka_dir, load_config, project_root
from abreka.models import ExperimentStatus, RemoteRunConfig
from abreka.persistence import (
    experiment_dir,
    load_experiment,
    normalize_exp_id,
    save_experiment,
    update_index_entry,
)
from abreka.run_executor import RemoteRunExecutor
from abreka.runner import AutonomousRunner
from abreka.runtime_costs import record_stage_runtime


@click.group("remote")
def remote_group():
    """Inspect and recover remote runs for an experiment."""


def _load_remote_run_metadata(metadata_path: Path) -> dict:
    if not metadata_path.is_file():
        raise click.ClickException(
            f"No remote run metadata found at {metadata_path}. "
            "This experiment does not appear to have a recoverable remote run."
        )
    try:
        return json.loads(metadata_path.read_text())
    except json.JSONDecodeError as e:
        raise click.ClickException(f"Invalid remote run metadata at {metadata_path}: {e}") from e


def _build_recovery_executor(abreka_dir: Path, metadata: dict, console: Console) -> RemoteRunExecutor:
    config = load_config(abreka_dir)
    host = metadata.get("host") or config.run.remote.host
    project_dir = metadata.get("remote_project_dir") or config.run.remote.project_dir
    if not host or not project_dir:
        raise click.ClickException(
            "Remote recovery metadata is missing host/project_dir and config.toml does not supply them."
        )

    remote = RemoteRunConfig(
        host=host,
        project_dir=project_dir,
        bootstrap_command=metadata.get("bootstrap_command", config.run.remote.bootstrap_command),
        poll_seconds=int(metadata.get("poll_seconds", config.run.remote.poll_seconds or 5)),
        startup_command=metadata.get("startup_command", config.run.remote.startup_command),
        teardown_command=metadata.get("teardown_command", config.run.remote.teardown_command),
    )
    return RemoteRunExecutor(project_root=project_root(abreka_dir), remote=remote, console=console)


def _resolve_remote_paths(executor: RemoteRunExecutor, metadata: dict, exp_id: str) -> dict[str, str]:
    if isinstance(metadata.get("remote_paths"), dict):
        return metadata["remote_paths"]
    run_relative_path = metadata.get("run_relative_path") or f"abreka/experiments/{exp_id}/run.py"
    return executor._build_remote_paths(Path(run_relative_path))


def _load_recovery_context_for_abreka_dir(abreka_dir: Path, exp_id: str, console: Console | None = None) -> dict:
    exp_id = normalize_exp_id(exp_id)
    exp_path = experiment_dir(abreka_dir, exp_id)
    metadata_path = exp_path / "artifacts" / "remote_run.json"
    metadata = _load_remote_run_metadata(metadata_path)
    console = console or Console()
    executor = _build_recovery_executor(abreka_dir, metadata, console)
    remote_paths = _resolve_remote_paths(executor, metadata, exp_id)
    return {
        "abreka_dir": abreka_dir,
        "exp_id": exp_id,
        "exp_path": exp_path,
        "exp": load_experiment(abreka_dir, exp_id),
        "metadata": metadata,
        "metadata_path": metadata_path,
        "console": console,
        "executor": executor,
        "remote_paths": remote_paths,
    }


def _load_recovery_context(exp_id: str) -> dict:
    return _load_recovery_context_for_abreka_dir(find_abreka_dir(), exp_id)


def _iter_remote_run_entries(abreka_dir: Path) -> list[dict]:
    entries: list[dict] = []
    exp_root = abreka_dir / "experiments"
    if not exp_root.is_dir():
        return entries

    for exp_path in sorted(p for p in exp_root.iterdir() if p.is_dir()):
        metadata_path = exp_path / "artifacts" / "remote_run.json"
        if not metadata_path.is_file():
            continue
        metadata = _load_remote_run_metadata(metadata_path)
        exp_id = exp_path.name
        entries.append({
            "exp_id": exp_id,
            "exp_path": exp_path,
            "metadata_path": metadata_path,
            "metadata": metadata,
            "exp": load_experiment(abreka_dir, exp_id),
        })
    return entries


def _is_recoverable_remote_run(metadata: dict) -> bool:
    return metadata.get("state") in {"starting", "running", "interrupted"}


def find_recoverable_remote_run_ids(abreka_dir: Path) -> list[str]:
    return [entry["exp_id"] for entry in _iter_remote_run_entries(abreka_dir) if _is_recoverable_remote_run(entry["metadata"])]


def resume_remote_run(
    abreka_dir: Path,
    exp_id: str,
    *,
    no_evaluate: bool = False,
    debug: bool = False,
    console: Console | None = None,
) -> str:
    recovery = _load_recovery_context_for_abreka_dir(abreka_dir, exp_id, console=console)
    exp_id = recovery["exp_id"]

    started_at = datetime.now(timezone.utc)
    started_monotonic = time.monotonic()
    stage_status = "ok"

    try:
        try:
            result = recovery["executor"].resume_recovery(
                output_path=recovery["exp_path"] / "artifacts" / "run_output.txt",
                metadata_path=recovery["metadata_path"],
                remote_paths=recovery["remote_paths"],
            )
        except Exception as e:
            stage_status = "error"
            raise click.ClickException(str(e)) from e

        if result.timed_out:
            stage_status = "timed_out"
        elif result.returncode != 0:
            stage_status = "error"

        runner = AutonomousRunner(abreka_dir=abreka_dir, debug=debug, console=recovery["console"])
        runner._mark_experiment_running(exp_id)
        try:
            runner._apply_run_result(
                exp_id,
                result,
                load_config(abreka_dir).run.timeout_minutes,
                recovery["exp_path"] / "artifacts" / "run_output.txt",
            )
        except Exception as e:
            if stage_status == "ok":
                stage_status = "error"
            raise click.ClickException(str(e)) from e

        if not no_evaluate:
            try:
                runner._evaluate(exp_id)
            except Exception as e:
                raise click.ClickException(str(e)) from e

        return exp_id
    finally:
        record_stage_runtime(
            abreka_dir,
            exp_id,
            stage="run",
            elapsed_seconds=time.monotonic() - started_monotonic,
            status=stage_status,
            started_at=started_at,
            finished_at=datetime.now(timezone.utc),
        )


@remote_group.command("list")
def remote_list_cmd():
    """List experiments with remote-run recovery metadata."""
    abreka_dir = find_abreka_dir()
    entries = _iter_remote_run_entries(abreka_dir)
    if not entries:
        click.echo("No remote runs found.")
        return

    for entry in entries:
        metadata = entry["metadata"]
        exp = entry["exp"]
        click.echo(
            f"{entry['exp_id']}  local={exp.status.value}  remote={metadata.get('state', 'unknown')}  "
            f"host={metadata.get('host', '(unknown)')}"
        )


@remote_group.command("status")
@click.argument("exp_id")
def remote_status_cmd(exp_id: str):
    """Show current recovery state for a remote run."""
    ctx = _load_recovery_context(exp_id)
    exp = ctx["exp"]
    metadata = ctx["metadata"]
    output_path = ctx["exp_path"] / "artifacts" / "run_output.txt"
    local_log_size = output_path.stat().st_size if output_path.is_file() else 0

    info = ctx["executor"].inspect_recovery(
        remote_paths=ctx["remote_paths"],
        remote_pid=metadata.get("remote_pid", ""),
    )

    click.echo(f"Experiment: {ctx['exp_id']}")
    click.echo(f"Local status: {exp.status.value}")
    click.echo(f"Remote state: {metadata.get('state', 'unknown')}")
    click.echo(f"Remote host: {metadata.get('host', '(unknown)')}")
    click.echo(f"Remote pid: {metadata.get('remote_pid', '(none)')}")
    click.echo(f"Remote project dir: {metadata.get('remote_project_dir', '(unknown)')}")
    click.echo(f"Local log: {output_path} ({local_log_size} bytes)")
    click.echo(f"Recovery file: {ctx['metadata_path']}")

    if info.get("remote_reachable"):
        process_state = "alive" if info.get("pid_alive") else "not running"
        click.echo(f"Remote process: {process_state}")
        status = info.get("status")
        if status is None:
            click.echo("Remote status file: not present")
        else:
            click.echo(f"Remote status file: {json.dumps(status)}")
        click.echo(f"Remote log size: {info.get('remote_log_size', 0)} bytes")
    else:
        click.echo("Remote process: unreachable")
        if info.get("error"):
            click.echo(f"Remote error: {info['error']}")


@remote_group.command("sync")
@click.argument("exp_id")
def remote_sync_cmd(exp_id: str):
    """Fetch the latest remote artifacts without waiting for completion."""
    ctx = _load_recovery_context(exp_id)
    try:
        ctx["executor"].sync_recovery_artifacts(
            remote_paths=ctx["remote_paths"],
            local_artifacts_dir=ctx["exp_path"] / "artifacts",
            metadata_path=ctx["metadata_path"],
        )
    except Exception as e:
        raise click.ClickException(str(e)) from e

    click.echo(f"Synced remote artifacts for experiment {ctx['exp_id']}")


@remote_group.command("resume")
@click.argument("exp_id")
@click.option("--no-evaluate", is_flag=True, help="Do not run the evaluate step after recovery.")
@click.pass_context
def remote_resume_cmd(ctx, exp_id: str, no_evaluate: bool):
    """Resume watching/finalizing an interrupted remote run."""
    abreka_dir = find_abreka_dir()
    debug = ctx.obj.get("debug", False) if ctx.obj else False
    recovered_id = resume_remote_run(abreka_dir, exp_id, no_evaluate=no_evaluate, debug=debug)
    click.echo(f"Recovered remote run for experiment {recovered_id}")


@remote_group.command("stop")
@click.argument("exp_id")
@click.option("--kill", is_flag=True, help="Send SIGKILL instead of SIGTERM.")
def remote_stop_cmd(exp_id: str, kill: bool):
    """Stop a remote run and mark the experiment failed."""
    ctx = _load_recovery_context(exp_id)

    try:
        ctx["executor"].stop_recovery(
            remote_pid=ctx["metadata"].get("remote_pid", ""),
            remote_paths=ctx["remote_paths"],
            local_artifacts_dir=ctx["exp_path"] / "artifacts",
            metadata_path=ctx["metadata_path"],
            kill=kill,
        )
    except Exception as e:
        raise click.ClickException(str(e)) from e

    exp = load_experiment(ctx["abreka_dir"], ctx["exp_id"])
    exp.status = ExperimentStatus.FAILED
    exp.failure_reason = "Remote run stopped by user" if not kill else "Remote run stopped by user (SIGKILL)"
    save_experiment(exp, ctx["abreka_dir"])
    update_index_entry(ctx["abreka_dir"], exp)

    click.echo(f"Stopped remote run for experiment {ctx['exp_id']}")
