from dataclasses import dataclass
from functools import cached_property
from typing import overload
from lsprotocol.types import Position, Range

from parsimonious.nodes import Node


@dataclass(frozen=True)
class Locator:
    raw: list[str]

    def _locate_pos(self, index: int) -> Position:
        counter = 0
        line = 0
        for ix, raw_line in enumerate(self.raw):
            if counter + len(raw_line) > index:
                line = ix
                break
            else:
                counter += len(raw_line)

        return Position(line, index - counter)

    @overload
    def locate(self, index: int) -> Position:
        ...

    @overload
    def locate(self, index: Node) -> Range:
        ...

    def locate(self, index: int | Node) -> Position | Range:
        if isinstance(index, int):
            return self._locate_pos(index)

        start = self._locate_pos(index.start)
        end = self._locate_pos(index.end)

        return Range(start, end)