from __future__ import annotations

import os
import shutil
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from .db import Db


def _parse_iso_z(ts: str) -> datetime:
    # Stored as "YYYY-MM-DDTHH:MM:SSZ" (naive UTC). Parse into aware UTC.
    return datetime.fromisoformat(ts.replace("Z", "+00:00")).astimezone(timezone.utc)


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


@dataclass
class JobsJanitor:
    db: Db
    user_root: str
    trash_after_days: int = 3
    purge_after_days: int = 7
    interval_s: int = 3600

    def __post_init__(self) -> None:
        if int(self.trash_after_days) < 0 or int(self.purge_after_days) < 0:
            raise ValueError("retention days must be non-negative")
        if int(self.purge_after_days) and int(self.purge_after_days) < int(self.trash_after_days):
            raise ValueError("purge_after_days must be >= trash_after_days")

    def _job_dir(self, *, user_id: str, ray_submission_id: str) -> str:
        base = self.user_root.rstrip("/")
        return f"{base}/{user_id}/jobs/{ray_submission_id}"

    def _trash_dir(self, *, user_id: str, ray_submission_id: str) -> str:
        base = self.user_root.rstrip("/")
        return f"{base}/{user_id}/trash/jobs/{ray_submission_id}"

    def tick_once(self, *, now: datetime | None = None, limit: int = 2000) -> None:
        if int(self.trash_after_days) <= 0 and int(self.purge_after_days) <= 0:
            return

        now_dt = now or datetime.now(timezone.utc)
        move_cutoff = now_dt - timedelta(days=int(self.trash_after_days))
        move_cutoff_iso = _iso_z(move_cutoff)

        rows = self.db.list_ended_attempts_before(end_time_le=move_cutoff_iso, limit=int(limit))
        for r in rows:
            user_id = str(r.get("user_id") or "").strip()
            sid = str(r.get("ray_submission_id") or "").strip()
            end_time = str(r.get("end_time") or "").strip()
            if not user_id or not sid or not end_time:
                continue

            try:
                ended_at = _parse_iso_z(end_time)
            except Exception:
                continue

            age_days = (now_dt - ended_at).total_seconds() / 86400.0
            src = self._job_dir(user_id=user_id, ray_submission_id=sid)
            dst = self._trash_dir(user_id=user_id, ray_submission_id=sid)
            dst_parent = os.path.dirname(dst)

            # Between trash and purge: ensure in trash.
            if age_days >= float(self.trash_after_days) and age_days < float(self.purge_after_days or 10**9):
                if os.path.exists(src) and not os.path.exists(dst):
                    os.makedirs(dst_parent, exist_ok=True)
                    try:
                        shutil.move(src, dst)
                    except Exception:
                        pass
                continue

            # Purge: move to trash (if still in jobs) then delete from trash.
            if int(self.purge_after_days) > 0 and age_days >= float(self.purge_after_days):
                if os.path.exists(src) and not os.path.exists(dst):
                    os.makedirs(dst_parent, exist_ok=True)
                    try:
                        shutil.move(src, dst)
                    except Exception:
                        pass
                if os.path.exists(dst):
                    try:
                        shutil.rmtree(dst)
                    except Exception:
                        pass

    def run_forever(self, stop_flag: object) -> None:
        try:
            is_set = stop_flag.is_set  # type: ignore[attr-defined]
        except Exception:
            raise ValueError("stop_flag must be a threading.Event-like object with is_set()")

        while not is_set():
            try:
                self.tick_once()
            except Exception:
                pass
            time.sleep(max(1, int(self.interval_s)))

