import os

from conftest import BUCKET_NAME, PREFIX, STS_CREDS

from vcm_file_uploader.segmenter import LogSegmenter, Segment
from vcm_file_uploader.session import STSUploadSession
from vcm_file_uploader.state import JsonlStore


def _make_session(**overrides):
    kwargs = {**STS_CREDS, "bucket": BUCKET_NAME, "prefix": PREFIX}
    kwargs.update(overrides)
    return STSUploadSession(**kwargs)


def _make_segmenter(tmp_path):
    store = JsonlStore(tmp_path / "state" / "ascii_segments.jsonl", key_field="segment_key")
    return LogSegmenter(store), store


class TestSegment:
    def test_segment_key_suffix(self):
        seg = Segment(path="logs/amber.mdout", offset=0, length=1024, segment_index=3)
        assert seg.segment_key_suffix == "logs/amber.mdout.seg0003"

    def test_segment_key_suffix_zero(self):
        seg = Segment(path="output.log", offset=0, length=100, segment_index=0)
        assert seg.segment_key_suffix == "output.log.seg0000"


class TestGetNewSegments:
    def test_new_file_full_segment(self, tmp_path):
        segmenter, _ = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("line 1\nline 2\nline 3\n")

        segments = segmenter.get_new_segments(str(log))

        assert len(segments) == 1
        assert segments[0].offset == 0
        assert segments[0].length == log.stat().st_size
        assert segments[0].segment_index == 0

    def test_no_new_data(self, tmp_path):
        segmenter, store = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("existing data\n")

        # Simulate already uploaded
        store.append({
            "segment_key": str(log),
            "last_uploaded_offset": log.stat().st_size,
            "segment_count": 1,
        })

        segments = segmenter.get_new_segments(str(log))
        assert len(segments) == 0

    def test_incremental_new_data(self, tmp_path):
        segmenter, store = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("line 1\n")
        initial_size = log.stat().st_size

        # Mark initial data as uploaded
        store.append({
            "segment_key": str(log),
            "last_uploaded_offset": initial_size,
            "segment_count": 1,
        })

        # File grows
        with open(log, "a") as f:
            f.write("line 2\nline 3\n")

        segments = segmenter.get_new_segments(str(log))
        assert len(segments) == 1
        assert segments[0].offset == initial_size
        assert segments[0].length == log.stat().st_size - initial_size
        assert segments[0].segment_index == 1

    def test_missing_file(self, tmp_path):
        segmenter, _ = _make_segmenter(tmp_path)
        segments = segmenter.get_new_segments(str(tmp_path / "missing.log"))
        assert len(segments) == 0

    def test_chunked_segments(self, tmp_path):
        segmenter, _ = _make_segmenter(tmp_path)
        log = tmp_path / "big.log"
        log.write_bytes(b"x" * 1000)

        segments = segmenter.get_new_segments(str(log), segment_size=300)

        assert len(segments) == 4  # 300 + 300 + 300 + 100
        assert segments[0].offset == 0
        assert segments[0].length == 300
        assert segments[1].offset == 300
        assert segments[1].length == 300
        assert segments[3].offset == 900
        assert segments[3].length == 100


class TestGetLastOffset:
    def test_no_state(self, tmp_path):
        segmenter, _ = _make_segmenter(tmp_path)
        assert segmenter.get_last_offset("some/file.log") == 0

    def test_with_state(self, tmp_path):
        segmenter, store = _make_segmenter(tmp_path)
        store.append({
            "segment_key": "some/file.log",
            "last_uploaded_offset": 4096,
            "segment_count": 2,
        })
        assert segmenter.get_last_offset("some/file.log") == 4096


class TestExtractSegment:
    def test_extract_full_file(self, tmp_path):
        segmenter, _ = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("hello world\n")

        seg = Segment(path=str(log), offset=0, length=12, segment_index=0)
        tmp_file = segmenter.extract_segment(seg)

        try:
            with open(tmp_file, "rb") as f:
                assert f.read() == b"hello world\n"
        finally:
            os.unlink(tmp_file)

    def test_extract_partial(self, tmp_path):
        segmenter, _ = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("aaabbbccc")

        seg = Segment(path=str(log), offset=3, length=3, segment_index=1)
        tmp_file = segmenter.extract_segment(seg)

        try:
            with open(tmp_file, "rb") as f:
                assert f.read() == b"bbb"
        finally:
            os.unlink(tmp_file)


class TestUploadSegment:
    def test_upload_and_update_state(self, mock_s3, tmp_path):
        bucket_name, client = mock_s3
        segmenter, store = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("line 1\nline 2\n")

        seg = Segment(
            path=str(log), offset=0, length=log.stat().st_size, segment_index=0,
        )
        s3_key = PREFIX + "logs/app.log.seg0000"

        with _make_session() as session:
            result = segmenter.upload_segment(seg, session, s3_key)

        assert result.success is True

        # Check S3
        obj = client.get_object(Bucket=bucket_name, Key=s3_key)
        assert obj["Body"].read() == b"line 1\nline 2\n"

        # Check state updated
        state = store.get(str(log))
        assert state is not None
        assert state["last_uploaded_offset"] == log.stat().st_size
        assert state["segment_count"] == 1

    def test_incremental_upload(self, mock_s3, tmp_path):
        bucket_name, client = mock_s3
        segmenter, store = _make_segmenter(tmp_path)
        log = tmp_path / "app.log"
        log.write_text("first\n")
        first_size = log.stat().st_size

        # Upload first segment
        seg0 = Segment(path=str(log), offset=0, length=first_size, segment_index=0)
        with _make_session() as session:
            segmenter.upload_segment(seg0, session, PREFIX + "logs/app.log.seg0000")

        # File grows
        with open(log, "a") as f:
            f.write("second\n")

        # Get and upload new segment
        segments = segmenter.get_new_segments(str(log))
        assert len(segments) == 1
        assert segments[0].segment_index == 1

        with _make_session() as session:
            result = segmenter.upload_segment(
                segments[0], session, PREFIX + "logs/app.log.seg0001"
            )

        assert result.success is True

        # Verify both segments in S3
        obj0 = client.get_object(Bucket=bucket_name, Key=PREFIX + "logs/app.log.seg0000")
        assert obj0["Body"].read() == b"first\n"
        obj1 = client.get_object(Bucket=bucket_name, Key=PREFIX + "logs/app.log.seg0001")
        assert obj1["Body"].read() == b"second\n"

        # State should reflect final offset
        state = store.get(str(log))
        assert state["last_uploaded_offset"] == log.stat().st_size
        assert state["segment_count"] == 2
