from typing import List
from dataclasses import dataclass


@dataclass
class Segment:
    segment_id: str
    start_line: int
    end_line: int
    start_byte: int
    end_byte: int
    content: str


class SourceSegmenter:
    def __init__(self, code: str):
        self.code = code
        self.lines = code.split('\n')

    @staticmethod
    def segment(code: str) -> List[Segment]:
        segmenter = SourceSegmenter(code)
        return segmenter._find_segments()

    def _find_segments(self) -> List[Segment]:
        boundaries = self._find_boundaries()
        segments = []

        for i, (start, end) in enumerate(boundaries):
            segment_id = f"seg_{i}"
            start_byte = self._line_offset(start)
            end_byte = self._line_offset(end)
            content = '\n'.join(self.lines[start:end])

            segments.append(Segment(
                segment_id=segment_id,
                start_line=start,
                end_line=end,
                start_byte=start_byte,
                end_byte=end_byte,
                content=content
            ))

        return segments

    def _line_offset(self, line_num: int) -> int:
        total = 0
        for i in range(min(line_num, len(self.lines))):
            total += len(self.lines[i]) + 1
        return total

    def _find_boundaries(self) -> List[tuple]:
        boundaries = []
        start = 0

        for i, line in enumerate(self.lines):
            stripped = line.lstrip()

            if stripped.startswith('def ') or stripped.startswith('class '):
                if start < i and any(self.lines[start:i]):
                    boundaries.append((start, i))
                start = i

        if start < len(self.lines):
            boundaries.append((start, len(self.lines)))

        return boundaries if boundaries else [(0, len(self.lines))]