"""abreka run — autonomous experiment loop."""

from __future__ import annotations

from dataclasses import dataclass

import click

from abreka.cli.exp.remote import find_recoverable_remote_run_ids, resume_remote_run
from abreka.config import find_abreka_dir, load_config, load_index
from abreka.models import ExperimentStatus
from abreka.runner import AutonomousRunner, parse_duration
from abreka.stage_sessions import stage_session_path


def _resolve_explore_probability(
    explore: bool,
    explore_prob: float | None,
    config_prob: float,
) -> float | None:
    """Resolve explore probability from CLI flags and config.

    Returns None if no exploration, or a float in (0, 1].
    Raises ValueError on invalid combinations.
    """
    if explore and explore_prob is not None:
        raise ValueError("--explore and --explore-prob are mutually exclusive.")
    if explore_prob is not None:
        if not (0.0 < explore_prob <= 1.0):
            raise ValueError("--explore-prob must be in (0.0, 1.0].")
        return explore_prob
    if explore:
        return 1.0
    if config_prob > 0:
        if config_prob > 1.0:
            raise ValueError("explore_probability in config.toml must be in [0.0, 1.0].")
        return config_prob
    return None


@dataclass(frozen=True)
class RecoverableCandidate:
    kind: str
    exp_id: str
    stage: str

    def label(self) -> str:
        return f"{self.exp_id} {self.stage}"


def _find_recoverable_remote_run_ids(abreka_dir) -> list[str]:
    return find_recoverable_remote_run_ids(abreka_dir)


def _find_recoverable_local_stage_candidates(abreka_dir) -> list[RecoverableCandidate]:
    index = load_index(abreka_dir)
    stage_by_status = {
        ExperimentStatus.PROPOSING.value: "propose",
        ExperimentStatus.IMPLEMENTING.value: "implement",
        ExperimentStatus.TESTING.value: "test",
        ExperimentStatus.RUNNING.value: "evaluate",
    }
    found: list[RecoverableCandidate] = []
    for exp_id in sorted(index):
        stage = stage_by_status.get(index[exp_id].get("status", ""))
        if not stage:
            continue
        if not stage_session_path(abreka_dir, exp_id, stage).is_file():
            continue
        found.append(RecoverableCandidate("local_stage", exp_id, stage))
    return found


def _find_recoverable_candidates(abreka_dir) -> list[RecoverableCandidate]:
    candidates = _find_recoverable_local_stage_candidates(abreka_dir)
    candidates.extend(
        RecoverableCandidate("remote_run", exp_id, "remote-run")
        for exp_id in _find_recoverable_remote_run_ids(abreka_dir)
    )
    return candidates


def _recover_local_stage(ctx, runner: AutonomousRunner, candidate: RecoverableCandidate) -> None:
    exp_id = candidate.exp_id
    stage = candidate.stage

    if stage == "propose":
        from abreka.cli.agent.propose import propose_cmd

        propose_cmd.main(args=["--resume", exp_id], standalone_mode=False, obj=ctx.obj)
        runner._implement(exp_id)
        runner._test(exp_id)
        runner._run_experiment(exp_id)
        runner._evaluate(exp_id)
        return

    if stage == "implement":
        from abreka.cli.agent.implement import implement_cmd

        implement_cmd.main(args=[exp_id, "--resume"], standalone_mode=False, obj=ctx.obj)
        runner._test(exp_id)
        runner._run_experiment(exp_id)
        runner._evaluate(exp_id)
        return

    if stage == "test":
        from abreka.cli.agent.test_cmd import test_cmd

        test_cmd.main(args=[exp_id, "--resume"], standalone_mode=False, obj=ctx.obj)
        runner._run_experiment(exp_id)
        runner._evaluate(exp_id)
        return

    if stage == "evaluate":
        from abreka.cli.agent.evaluate import evaluate_cmd

        evaluate_cmd.main(args=[exp_id, "--resume"], standalone_mode=False, obj=ctx.obj)
        return

    raise click.ClickException(f"Unsupported recoverable stage: {stage}")


def _recover_candidate(ctx, abreka_dir, runner: AutonomousRunner, candidate: RecoverableCandidate, *, debug: bool) -> str:
    if candidate.kind == "remote_run":
        return resume_remote_run(abreka_dir, candidate.exp_id, debug=debug)
    _recover_local_stage(ctx, runner, candidate)
    return candidate.exp_id


@click.command("run")
@click.argument("duration", required=False, default=None)
@click.option("--step", is_flag=True, help="Run a single iteration then stop.")
@click.option("--recover", is_flag=True, help="Recover a single interrupted stage or remote run before starting new work.")
@click.option("--explore", is_flag=True, help="Use exploration prompt (always explore).")
@click.option(
    "--explore-prob",
    type=float,
    default=None,
    help="Probability (0.0-1.0] of exploring each iteration.",
)
@click.option("--focus", "focus_root", default=None, help="Focus on subtree rooted at this experiment ID.")
@click.option(
    "--guidance",
    default=None,
    is_flag=False,
    flag_value="__interactive__",
    help="Guidance for the proposal agent. Pass a string, or use bare --guidance for interactive input.",
)
@click.option(
    "--guidance-file",
    default=None,
    type=click.Path(exists=True),
    help="Path to a file containing guidance text.",
)
@click.pass_context
def run_cmd(ctx, duration: str | None, step: bool, recover: bool, explore: bool, explore_prob: float | None, focus_root: str | None, guidance: str | None, guidance_file: str | None):
    """Run the autonomous experiment loop.

    DURATION can be e.g. '24h', '2h', '30m'. Use --step for a single iteration.
    """
    if not duration and not step and not recover:
        raise click.ClickException("Specify a duration (e.g. '24h'), use --step, or pass --recover.")

    debug = ctx.obj.get("debug", False) if ctx.obj else False
    abreka_dir = find_abreka_dir()
    config = load_config(abreka_dir)

    if focus_root:
        from abreka.config import load_index
        index = load_index(abreka_dir)
        if focus_root not in index:
            raise click.ClickException(f"Experiment {focus_root!r} not found.")

    if step:
        duration_seconds = None
    elif duration:
        duration_seconds = parse_duration(duration)
    else:
        duration_seconds = None

    try:
        explore_probability = _resolve_explore_probability(
            explore, explore_prob, config.experiment.explore_probability,
        )
    except ValueError as e:
        raise click.ClickException(str(e))

    # Resolve guidance
    from abreka.cli.agent.propose import _read_guidance_interactive, _resolve_guidance

    if guidance == "__interactive__":
        import click as _click
        _click.echo("Enter guidance (end with EOF on its own line):")
        guidance = _read_guidance_interactive()
    resolved_guidance = _resolve_guidance(guidance, guidance_file)

    runner = AutonomousRunner(
        abreka_dir=abreka_dir,
        duration_seconds=duration_seconds,
        explore_probability=explore_probability,
        focus_root=focus_root,
        guidance=resolved_guidance,
        debug=debug,
    )

    if recover:
        candidates = _find_recoverable_candidates(abreka_dir)
        if not candidates:
            raise click.ClickException("No recoverable runs found.")
        if len(candidates) > 1:
            choices = ", ".join(c.label() for c in candidates)
            raise click.ClickException(
                "Multiple recoverable runs found: "
                f"{choices}. Resume one explicitly with the stage command or 'abreka exp remote resume <id>'."
            )
        recovered = candidates[0]
        recovered_id = _recover_candidate(ctx, abreka_dir, runner, recovered, debug=debug)
        click.echo(f"Recovered {recovered.stage} for experiment {recovered_id}")
        if not duration and not step:
            return

    runner.run()
