from typing import NamedTuple, cast

from imandra.u.agents.code_logician.base import FormalizationState
from imandrax_api_models import ErrorMessage, Position

from ..util import filter, find, flat_map, map, scanl, snd


class ErrorSnippet(NamedTuple):
    file: str
    start_line: int
    start_col: int
    stop_line: int
    stop_col: int
    lines: list[str]


def full_iml_model(fstate: FormalizationState):
    """Get the full IML code, including the monolith IML redefining the dependencies
    modules and the original IML code.
    """
    if fstate.iml_model is None:
        raise ValueError('IML code is required')

    def fmt_module(module):
        indented = ['    ' + l for l in module.content.strip().splitlines()]
        lines = [f'module {module.name} = struct', *indented, 'end', '']
        return str(module.relative_path), lines

    main = None, fstate.iml_model.splitlines()
    return [fmt_module(d.iml_module) for d in fstate.dependency] + [main]


def join_sections(sections):
    return flat_map(snd, sections)


def find_line(sections, line):
    cum_sums = list(scanl(lambda x, y: x + len(snd(y)), 0, sections))
    n = cast(int, find(lambda i: line < cum_sums[i], range(len(cum_sums))))
    return line - cum_sums[n - 1], sections[n - 1]


def snippet_lines(
    lines: list[str], start: Position, end: Position, context_line: int = 2
):
    """Format a code snippet with highlighted error range."""
    start_line, start_col = start.line, start.col - 1
    end_line, end_col = end.line, end.col

    # Determine visible range
    visible_start = max(0, start_line - context_line - 1)
    visible_end = min(len(lines), end_line + context_line)

    # Calculate max line number width for alignment
    line_no_width = len(str(visible_end))
    underline_prefix = ' ' * (2 + line_no_width)

    def line(underline):
        return f'{underline_prefix} | {underline}'

    for i in range(visible_start, visible_end):
        line_no = i + 1

        # Mark lines within error range
        marker = '*' if start_line <= line_no <= end_line else ' '
        yield f'{marker} {line_no:{line_no_width}} | {lines[i]}'

        # Add underline/pointer
        if line_no == start_line == end_line:
            # Single line error
            yield line(' ' * start_col + '^' * (end_col - start_col))
        elif line_no == start_line:
            # Start of multi-line error
            yield line(' ' * start_col + '^' + '~' * (len(lines[i]) - start_col))
        elif line_no == end_line:
            # End of multi-line error
            yield line('~' * end_col)
        elif start_line < line_no < end_line:
            # Middle of multi-line error
            yield line('~' * len(lines[i]))


def describe_error(mod_path: str, sections, error_msg: ErrorMessage):
    def has_loc(loc):
        return loc.start is not None and loc.stop is not None

    def snippet(p) -> ErrorSnippet:
        n = p.stop.line - p.start.line
        l, (rel_path, lines) = find_line(sections, p.start.line - 1)
        ll, lines, path = (
            (l, lines[1:], rel_path) if rel_path else (l + 1, lines, mod_path)
        )
        start = Position(line=ll, col=p.start.col)
        stop = Position(line=ll + n, col=p.stop.col)
        snippet = list(snippet_lines(lines, start, stop))

        return ErrorSnippet(path, start.line, start.col, stop.line, stop.col, snippet)

    return map(snippet, filter(has_loc, error_msg.locs or []))
