"""
Chunked writer with checkpoint support for large sheet writes.

Used by sheets_io to stream writes and resume from checkpoints when requested.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator, List, Optional


@dataclass
class CheckpointState:
    completed_rows: int
    last_chunk: int


class ChunkedWriter:
    def __init__(self, checkpoint_file: Path, chunk_size: int = 500):
        self.checkpoint_file = checkpoint_file
        self.chunk_size = chunk_size

    def load_checkpoint(self) -> Optional[CheckpointState]:
        if not self.checkpoint_file.exists():
            return None
        try:
            data = json.loads(self.checkpoint_file.read_text(encoding="utf-8"))
            return CheckpointState(
                completed_rows=data.get("completed_rows", 0),
                last_chunk=data.get("last_chunk", 0),
            )
        except Exception:
            return None

    def save_checkpoint(self, completed_rows: int, last_chunk: int) -> None:
        self.checkpoint_file.parent.mkdir(parents=True, exist_ok=True)
        self.checkpoint_file.write_text(
            json.dumps(
                {
                    "completed_rows": completed_rows,
                    "last_chunk": last_chunk,
                },
                indent=2,
            ),
            encoding="utf-8",
        )

    def clear_checkpoint(self) -> None:
        if self.checkpoint_file.exists():
            self.checkpoint_file.unlink()

    def write_chunks(self, rows: Iterator[List[Any]], writer_func, resume: bool = False) -> int:
        """
        Write rows in chunks using provided writer_func(chunk_rows).
        If resume=True, skips rows already recorded in the checkpoint.
        Returns total rows processed (including any skipped rows on resume).
        """
        checkpoint = self.load_checkpoint() if resume else None
        completed = checkpoint.completed_rows if checkpoint else 0
        current_chunk = checkpoint.last_chunk + 1 if checkpoint else 1

        batch: List[List[Any]] = []
        total_processed = 0

        # Skip rows that were already written in a previous run
        skipped = 0
        if completed:
            for _ in range(completed):
                try:
                    next(rows)
                    skipped += 1
                except StopIteration:
                    break
            total_processed = skipped

        try:
            for row in rows:
                batch.append(row)
                if len(batch) >= self.chunk_size:
                    writer_func(batch)
                    total_processed += len(batch)
                    self.save_checkpoint(total_processed, current_chunk)
                    current_chunk += 1
                    batch = []

            if batch:
                writer_func(batch)
                total_processed += len(batch)
                self.save_checkpoint(total_processed, current_chunk)
        except Exception:
            # Leave checkpoint intact for resume
            raise

        self.clear_checkpoint()
        return total_processed
