from __future__ import annotations

import json
import shutil
from datetime import datetime
from pathlib import Path

from one_run.commands.cancel_cmd import _cancel_local, _cancel_slurm
from one_run.commands.common import load_run_paths
from one_run.utils.duration import parse_duration
from one_run.errors import OneRunError
from one_run.models import RunMetadata
from one_run.payloads import rm_cleanup_result_payload, rm_result_payload
from one_run.status_store import read_metadata, read_status


def remove_run(run_id: str, *, force: bool = False, json_output: bool = False) -> int:
    project_root = Path.cwd().resolve()
    paths = load_run_paths(project_root, run_id)
    metadata = RunMetadata.model_validate(read_metadata(paths.metadata))
    status = read_status(paths.status)

    state_before = status.state
    cancelled = False
    state_after = status.state

    if status.state not in {"success", "failed"}:
        if metadata.backend_type == "slurm":
            updated = _cancel_slurm(metadata, status, paths, force=force)
        else:
            updated = _cancel_local(paths, status, force=force)
        cancelled = True
        state_after = updated.state

    shutil.rmtree(paths.run_dir)

    if json_output:
        print(
            json.dumps(
                rm_result_payload(
                    run_id=run_id,
                    removed=True,
                    cancelled=cancelled,
                    force=force,
                    state_before=state_before,
                    state_after=state_after,
                )
            )
        )
    else:
        print(f"run_id={run_id}")
        print("removed=true")
        print(f"cancelled={str(cancelled).lower()}")
        print(f"force={str(force).lower()}")
    return 0


def _reference_time(status, metadata: dict) -> datetime | None:
    if status.ended_at:
        return status.ended_at
    if status.started_at:
        return status.started_at
    ts = metadata.get("timestamp")
    if not ts:
        return None
    try:
        return datetime.fromisoformat(str(ts))
    except ValueError:
        return None


def remove_runs(
    *,
    state: str | None = None,
    older_than: str | None = None,
    dry_run: bool = False,
    json_output: bool = False,
) -> int:
    project_root = Path.cwd().resolve()
    runs_root = project_root / "runs"
    if not runs_root.exists():
        if json_output:
            print(json.dumps(rm_cleanup_result_payload(items=[], count=0, dry_run=dry_run)))
        else:
            print("no runs found")
        return 0

    cutoff = None
    if older_than:
        try:
            cutoff = datetime.now() - parse_duration(older_than)
        except ValueError as exc:
            raise OneRunError(str(exc)) from exc

    selected: list[Path] = []
    for run_dir in sorted([p for p in runs_root.iterdir() if p.is_dir()], reverse=True):
        metadata_path = run_dir / "metadata.json"
        status_path = run_dir / "status.json"
        if not metadata_path.exists() or not status_path.exists():
            continue
        metadata = read_metadata(metadata_path)
        status_obj = read_status(status_path)
        if state and status_obj.state != state:
            continue
        ref_ts = _reference_time(status_obj, metadata)
        if cutoff and ref_ts and ref_ts >= cutoff:
            continue
        selected.append(run_dir)

    if not selected:
        if json_output:
            print(json.dumps(rm_cleanup_result_payload(items=[], count=0, dry_run=dry_run)))
        else:
            print("no runs matched filter")
        return 0

    items: list[dict[str, str]] = []
    for run_dir in selected:
        if dry_run:
            items.append({"run_id": run_dir.name, "action": "would_delete"})
            if not json_output:
                print(f"would_delete={run_dir.name}")
        else:
            shutil.rmtree(run_dir)
            items.append({"run_id": run_dir.name, "action": "deleted"})
            if not json_output:
                print(f"deleted={run_dir.name}")

    if json_output:
        print(json.dumps(rm_cleanup_result_payload(items=items, count=len(selected), dry_run=dry_run)))
    else:
        print(f"count={len(selected)}")
        print(f"dry_run={str(dry_run).lower()}")
    return 0
