"""Resume state tracking via .progress/ journal files."""

from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


class MetadataFetchState:
    """Tracks which metadata pages have been fetched for a query."""

    def __init__(self, progress_dir: Path):
        self._path = progress_dir / "metadata-fetch.json"
        self._state: dict[str, Any] = {}
        self._load()

    def _load(self) -> None:
        if self._path.exists():
            with open(self._path) as f:
                self._state = json.load(f)

    def _save(self) -> None:
        self._path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self._path.with_suffix(".tmp")
        with open(tmp, "w") as f:
            json.dump(self._state, f, indent=2)
        tmp.rename(self._path)

    def get_fetched_pages(self, query_hash: str) -> set[int]:
        if self._state.get("query_hash") != query_hash:
            return set()
        return set(self._state.get("fetched_pages", []))

    def mark_page(self, query_hash: str, page: int, total_pages: int) -> None:
        if self._state.get("query_hash") != query_hash:
            self._state = {
                "query_hash": query_hash,
                "total_pages": total_pages,
                "fetched_pages": [],
            }
        pages = set(self._state.get("fetched_pages", []))
        pages.add(page)
        self._state["fetched_pages"] = sorted(pages)
        self._state["last_updated"] = datetime.now(timezone.utc).isoformat()
        self._save()

    def is_complete(self, query_hash: str) -> bool:
        if self._state.get("query_hash") != query_hash:
            return False
        total = self._state.get("total_pages", 0)
        fetched = set(self._state.get("fetched_pages", []))
        return total > 0 and len(fetched) >= total


class DownloadState:
    """Tracks per-file download status."""

    def __init__(self, progress_dir: Path, suffix: str = ""):
        name = f"download-state{suffix}.json"
        self._path = progress_dir / name
        self._state: dict[str, dict[str, Any]] = {}
        self._load()

    def _load(self) -> None:
        if self._path.exists():
            with open(self._path) as f:
                self._state = json.load(f)

    def _save(self) -> None:
        self._path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self._path.with_suffix(".tmp")
        with open(tmp, "w") as f:
            json.dump(self._state, f)
        tmp.rename(self._path)

    def get_status(self, xc_id: str) -> str | None:
        entry = self._state.get(xc_id)
        return entry.get("status") if entry else None

    def mark_downloading(self, xc_id: str) -> None:
        self._state[xc_id] = {"status": "downloading"}
        self._save()

    def mark_verified(self, xc_id: str, sha256: str, size: int) -> None:
        self._state[xc_id] = {
            "status": "verified",
            "sha256": sha256,
            "file_size_bytes": size,
        }
        self._save()

    def mark_failed(self, xc_id: str, error: str, retries: int) -> None:
        self._state[xc_id] = {
            "status": "failed",
            "error": error,
            "retries": retries,
        }
        self._save()

    def get_pending_ids(self, all_ids: set[str], retry_failed: bool = True) -> list[str]:
        """Return IDs that need downloading: not verified, plus optionally failed."""
        pending = []
        for xc_id in all_ids:
            status = self.get_status(xc_id)
            if status == "verified":
                continue
            if status == "failed" and not retry_failed:
                continue
            pending.append(xc_id)
        return pending
