"""LogSegmenter — handle growing ASCII log uploads via byte-range segments."""

from __future__ import annotations

import logging
import os
import tempfile
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from vcm_file_uploader.session import STSUploadSession

from vcm_file_uploader._types import UploadResult
from vcm_file_uploader.state import JsonlStore

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class Segment:
    """A byte range of a growing file to upload."""

    path: str
    offset: int
    length: int
    segment_index: int

    @property
    def segment_key_suffix(self) -> str:
        """S3 key suffix for this segment, e.g. 'logs/amber.mdout.seg0003'."""
        return f"{self.path}.seg{self.segment_index:04d}"


class LogSegmenter:
    """Tracks growing ASCII logs and uploads new segments.

    Each file is tracked by its last uploaded byte offset in JSONL state.
    On each call to get_new_segments(), returns byte ranges representing
    new data since the last upload.
    """

    def __init__(self, segments_store: JsonlStore) -> None:
        self._store = segments_store

    def _state_key(self, path: str) -> str:
        return path

    def get_last_offset(self, path: str) -> int:
        """Get the last uploaded byte offset for a file."""
        entry = self._store.get(self._state_key(path))
        if entry is None:
            return 0
        return int(entry.get("last_uploaded_offset", 0))

    def get_new_segments(self, path: str, segment_size: int = 0) -> list[Segment]:
        """Detect new bytes and return segments to upload.

        If segment_size is 0 (default), all new bytes are returned as
        a single segment. Otherwise, new bytes are split into chunks
        of segment_size.
        """
        try:
            current_size = os.path.getsize(path)
        except OSError:
            return []

        last_offset = self.get_last_offset(path)
        if current_size <= last_offset:
            return []

        new_bytes = current_size - last_offset

        # Determine segment count from state for indexing
        entry = self._store.get(self._state_key(path))
        next_index = entry.get("segment_count", 0) if entry else 0

        segments: list[Segment] = []
        offset = last_offset

        if segment_size <= 0:
            segments.append(Segment(
                path=path,
                offset=offset,
                length=new_bytes,
                segment_index=next_index,
            ))
        else:
            while offset < current_size:
                length = min(segment_size, current_size - offset)
                segments.append(Segment(
                    path=path,
                    offset=offset,
                    length=length,
                    segment_index=next_index,
                ))
                offset += length
                next_index += 1

        return segments

    def extract_segment(self, segment: Segment) -> str:
        """Extract a segment's bytes to a temporary file.

        Returns the path to the temp file. Caller is responsible for cleanup.
        """
        fd, tmp_path = tempfile.mkstemp(suffix=f".seg{segment.segment_index:04d}")
        try:
            with os.fdopen(fd, "wb") as f_out, open(segment.path, "rb") as f_in:
                f_in.seek(segment.offset)
                remaining = segment.length
                while remaining > 0:
                    chunk = f_in.read(min(remaining, 65536))
                    if not chunk:
                        break
                    f_out.write(chunk)
                    remaining -= len(chunk)
        except BaseException:
            os.unlink(tmp_path)
            raise
        return tmp_path

    def upload_segment(
        self,
        segment: Segment,
        session: STSUploadSession,
        s3_key: str,
    ) -> UploadResult:
        """Extract and upload a single segment, then update state."""
        tmp_path = self.extract_segment(segment)
        try:
            result = session.upload_file(tmp_path, s3_key)
        finally:
            os.unlink(tmp_path)

        if result.success:
            new_offset = segment.offset + segment.length
            self._store.append({
                "segment_key": self._state_key(segment.path),
                "path": segment.path,
                "last_uploaded_offset": new_offset,
                "segment_count": segment.segment_index + 1,
                "last_segment_key": s3_key,
            })
            logger.info(
                "Segment %s uploaded, offset now %d",
                s3_key, new_offset,
            )

        return result
