from __future__ import annotations

from datetime import datetime, timedelta, timezone
from pathlib import Path

from argus.service.db import Db
from argus.service.janitor import JobsJanitor


def _iso_z(dt: datetime) -> str:
    return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def _mk_job_dir(user_root: Path, user_id: str, sid: str) -> Path:
    p = user_root / user_id / "jobs" / sid
    p.mkdir(parents=True, exist_ok=True)
    (p / "marker.txt").write_text("x", encoding="utf-8")
    return p


def _mk_trash_dir(user_root: Path, user_id: str, sid: str) -> Path:
    p = user_root / user_id / "trash" / "jobs" / sid
    p.mkdir(parents=True, exist_ok=True)
    (p / "marker.txt").write_text("x", encoding="utf-8")
    return p


def test_janitor_moves_jobs_to_trash_after_3_days(tmp_path: Path) -> None:
    db_path = tmp_path / "mvp.sqlite3"
    user_root = tmp_path / "users"
    db = Db(str(db_path))
    db.init()

    task_id = "t1"
    user_id = "alice"
    sid = "sid-a01"
    db.create_task_v25(task_id=task_id, user_id=user_id, workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)

    now = datetime(2025, 1, 10, tzinfo=timezone.utc)
    ended = now - timedelta(days=4)
    db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="SUCCEEDED")

    src = _mk_job_dir(user_root, user_id, sid)
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
    jan.tick_once(now=now)

    assert not src.exists()
    dst = user_root / user_id / "trash" / "jobs" / sid
    assert dst.exists()
    assert (dst / "marker.txt").read_text(encoding="utf-8") == "x"


def test_janitor_purges_from_trash_after_7_days(tmp_path: Path) -> None:
    db_path = tmp_path / "mvp.sqlite3"
    user_root = tmp_path / "users"
    db = Db(str(db_path))
    db.init()

    task_id = "t2"
    user_id = "alice"
    sid = "sid-a01"
    db.create_task_v25(task_id=task_id, user_id=user_id, workload="ppo", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)

    now = datetime(2025, 1, 10, tzinfo=timezone.utc)
    ended = now - timedelta(days=8)
    db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="FAILED")

    _mk_trash_dir(user_root, user_id, sid)
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
    jan.tick_once(now=now)

    dst = user_root / user_id / "trash" / "jobs" / sid
    assert not dst.exists()


def test_janitor_does_not_touch_recent_jobs(tmp_path: Path) -> None:
    db_path = tmp_path / "mvp.sqlite3"
    user_root = tmp_path / "users"
    db = Db(str(db_path))
    db.init()

    task_id = "t3"
    user_id = "alice"
    sid = "sid-a01"
    db.create_task_v25(task_id=task_id, user_id=user_id, workload="grpo", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)

    now = datetime(2025, 1, 10, tzinfo=timezone.utc)
    ended = now - timedelta(days=1)
    db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="FAILED")

    src = _mk_job_dir(user_root, user_id, sid)
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
    jan.tick_once(now=now)

    assert src.exists()
    assert not (user_root / user_id / "trash" / "jobs" / sid).exists()


def test_janitor_skips_tasks_without_user_id(tmp_path: Path) -> None:
    db_path = tmp_path / "mvp.sqlite3"
    user_root = tmp_path / "users"
    db = Db(str(db_path))
    db.init()

    task_id = "legacy"
    sid = "sid-legacy"
    db.create_task(task_id=task_id, workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)

    now = datetime(2025, 1, 10, tzinfo=timezone.utc)
    ended = now - timedelta(days=10)
    db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="SUCCEEDED")

    # Even if someone created a matching directory under user_root, janitor shouldn't touch it because user_id is NULL.
    src = _mk_job_dir(user_root, "alice", sid)
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
    jan.tick_once(now=now)
    assert src.exists()


def test_janitor_validates_retention_days(tmp_path: Path) -> None:
    db = Db(str(tmp_path / "mvp.sqlite3"))
    db.init()
    try:
        JobsJanitor(db=db, user_root="/tmp", trash_after_days=-1, purge_after_days=7, interval_s=1)
        raise AssertionError("expected ValueError")
    except ValueError:
        pass
    try:
        JobsJanitor(db=db, user_root="/tmp", trash_after_days=3, purge_after_days=1, interval_s=1)
        raise AssertionError("expected ValueError")
    except ValueError:
        pass


def test_janitor_noop_when_disabled(tmp_path: Path) -> None:
    db = Db(str(tmp_path / "mvp.sqlite3"))
    db.init()
    user_root = tmp_path / "users"
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=0, purge_after_days=0, interval_s=1)
    jan.tick_once(now=datetime(2025, 1, 1, tzinfo=timezone.utc))


def test_janitor_handles_invalid_end_time_and_missing_fields(tmp_path: Path) -> None:
    db_path = tmp_path / "mvp.sqlite3"
    user_root = tmp_path / "users"
    db = Db(str(db_path))
    db.init()

    now = datetime(2025, 1, 10, tzinfo=timezone.utc)
    cutoff_ended = now - timedelta(days=10)

    # Missing end_time (empty string) => should be skipped.
    db.create_task_v25(task_id="t4", user_id="alice", workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id="t4", attempt_no=1, ray_submission_id="sid-empty")
    db.update_attempt(task_id="t4", attempt_no=1, end_time="", ray_status="SUCCEEDED")

    # Invalid ISO but still lexicographically <= cutoff => should be skipped.
    db.create_task_v25(task_id="t5", user_id="alice", workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id="t5", attempt_no=1, ray_submission_id="sid-bad")
    db.update_attempt(task_id="t5", attempt_no=1, end_time="2025-01-01T00:00:00ZZ", ray_status="FAILED")

    _mk_job_dir(user_root, "alice", "sid-empty")
    _mk_job_dir(user_root, "alice", "sid-bad")
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
    jan.tick_once(now=cutoff_ended + timedelta(days=10))

    assert (user_root / "alice" / "jobs" / "sid-empty").exists()
    assert (user_root / "alice" / "jobs" / "sid-bad").exists()


def test_janitor_purge_moves_from_jobs_then_deletes(tmp_path: Path, monkeypatch) -> None:
    db_path = tmp_path / "mvp.sqlite3"
    user_root = tmp_path / "users"
    db = Db(str(db_path))
    db.init()

    task_id = "t6"
    user_id = "alice"
    sid = "sid-a01"
    db.create_task_v25(task_id=task_id, user_id=user_id, workload="ppo", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
    db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)

    now = datetime(2025, 1, 10, tzinfo=timezone.utc)
    ended = now - timedelta(days=9)
    db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="SUCCEEDED")

    src = _mk_job_dir(user_root, user_id, sid)
    jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)

    jan.tick_once(now=now)
    assert not src.exists()
    assert not (user_root / user_id / "trash" / "jobs" / sid).exists()


def test_janitor_run_forever_requires_event_like(tmp_path: Path) -> None:
    db = Db(str(tmp_path / "mvp.sqlite3"))
    db.init()
    jan = JobsJanitor(db=db, user_root=str(tmp_path / "users"), trash_after_days=3, purge_after_days=7, interval_s=1)
    try:
        jan.run_forever(object())
        raise AssertionError("expected ValueError")
    except ValueError:
        pass


def test_janitor_run_forever_survives_tick_exceptions(tmp_path: Path, monkeypatch) -> None:
    db = Db(str(tmp_path / "mvp.sqlite3"))
    db.init()
    jan = JobsJanitor(db=db, user_root=str(tmp_path / "users"), trash_after_days=3, purge_after_days=7, interval_s=1)

    class Flag:
        def __init__(self) -> None:
            self.n = 0

        def is_set(self) -> bool:
            self.n += 1
            return self.n >= 2

    monkeypatch.setattr(jan, "tick_once", lambda **_: (_ for _ in ()).throw(RuntimeError("boom")))
    monkeypatch.setattr("argus.service.janitor.time.sleep", lambda *_: None)
    jan.run_forever(Flag())
