"""JSONL state management for resumable uploads and log segmentation."""

from __future__ import annotations

import json
import logging
import os
import tempfile
from pathlib import Path

logger = logging.getLogger(__name__)


class JsonlStore:
    """Append-only JSONL store with last-write-wins semantics.

    Each entry must contain a key field. On read, the latest entry
    for each key wins. Compaction deduplicates the file on disk.
    """

    def __init__(self, path: str | Path, key_field: str = "path") -> None:
        self.path = Path(path)
        self.key_field = key_field

    def append(self, entry: dict) -> None:
        """Append a single entry to the JSONL file."""
        if self.key_field not in entry:
            raise ValueError(f"Entry missing required key field: {self.key_field!r}")
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.path, "a") as f:
            f.write(json.dumps(entry, default=str) + "\n")

    def read_all(self) -> dict[str, dict]:
        """Read all entries, returning latest per key (last-write-wins)."""
        result: dict[str, dict] = {}
        if not self.path.exists():
            return result
        with open(self.path) as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except json.JSONDecodeError:
                    logger.warning("Skipping malformed JSONL line %d in %s", line_num, self.path)
                    continue
                key = entry.get(self.key_field)
                if key is not None:
                    result[key] = entry
        return result

    def get(self, key: str) -> dict | None:
        """Get the latest entry for a given key."""
        return self.read_all().get(key)

    def compact(self) -> int:
        """Deduplicate the file, keeping only the latest entry per key.

        Returns the number of entries after compaction.
        """
        entries = self.read_all()
        if not entries:
            if self.path.exists():
                self.path.unlink()
            return 0

        # Write to temp file then atomic rename
        self.path.parent.mkdir(parents=True, exist_ok=True)
        fd, tmp_path = tempfile.mkstemp(
            dir=self.path.parent, suffix=".jsonl.tmp"
        )
        try:
            with os.fdopen(fd, "w") as f:
                for entry in entries.values():
                    f.write(json.dumps(entry, default=str) + "\n")
            os.replace(tmp_path, self.path)
        except BaseException:
            os.unlink(tmp_path)
            raise

        return len(entries)


class StateManager:
    """Manages JSONL state files for the uploader.

    Provides access to three stores:
    - files: tracks uploaded files by path
    - segments: tracks ASCII log segments by segment_key
    - runs: tracks upload runs by run_id
    """

    def __init__(self, state_dir: str | Path) -> None:
        self.state_dir = Path(state_dir)
        self.files = JsonlStore(self.state_dir / "files.jsonl", key_field="path")
        self.segments = JsonlStore(self.state_dir / "ascii_segments.jsonl", key_field="segment_key")
        self.runs = JsonlStore(self.state_dir / "runs.jsonl", key_field="run_id")

    def compact_all(self) -> dict[str, int]:
        """Compact all state files. Returns counts per store."""
        return {
            "files": self.files.compact(),
            "segments": self.segments.compact(),
            "runs": self.runs.compact(),
        }
