from typing import Optional

from lsprotocol.types import Position

from .aws_data import AWSRefName, AWSRefSource, Tree
from .decode.extractors import (
    CompositeExtractor,
    Extractor,
    KeyExtractor,
    LogicalIdExtractor,
    ParameterExtractor,
)
from .decode.position import PositionLink, Spanning

REF_EXTRACTOR = KeyExtractor[AWSRefName]("Ref", lambda s: AWSRefName(value=s))
REF_SRC_EXTRACTOR = CompositeExtractor[AWSRefSource](
    LogicalIdExtractor(), ParameterExtractor()
)


def resolve_ref(
    position: Position,
    template_data: Tree,
    ref_extractor: Extractor[AWSRefName] = REF_EXTRACTOR,
    ref_src_extractor: Extractor[AWSRefSource] = REF_SRC_EXTRACTOR,
) -> Optional[PositionLink[AWSRefSource, AWSRefName]]:
    """Attempt to resolve the source and documentation of a ref at position.

    Parameters
    ----------
    position : Position
        The position to look at in the template
    template_data : Tree
        The decoded template
    extractor: Extractor[AWSRefSource]
        The extractor to use to extract the source of refs

    Returns
    -------
    Optional[PositionLink[AWSRefSource, AWSRefName]]
        A spanning object containing the ref source if it was found, else None
    """
    ref_lookup = ref_extractor.extract(template_data)
    ref_span = ref_lookup.at(position.line, position.character)
    if ref_span:
        ref_src_lookup = ref_src_extractor.extract(template_data)
        for ref_src, position_list in ref_src_lookup.items():
            if ref_src.logical_name == ref_span.value.value and position_list:
                line, char, _ = position_list[0]
                src_span = Spanning[AWSRefSource](
                    value=ref_src, line=line, char=char, span=len(ref_src.logical_name)
                )
                return PositionLink[AWSRefSource, AWSRefName](
                    source_span=src_span, target_span=ref_span
                )
    return None
