from dataclasses import dataclass
from typing import List, Optional


def middle_truncate(text: str, limit: int, newline_marker: bool) -> tuple[str, bool]:
    if len(text) <= limit:
        return text, False

    head_len = limit // 2
    tail_len = limit - head_len
    suffix = text[-tail_len:]
    marker = f"...[TRUNCATED {len(text) - limit} chars]..."
    if newline_marker:
        marker = f"\n{marker}\n"
    return text[:head_len] + marker + suffix, True


@dataclass
class TruncationResult:
    text: str
    orig_text: str
    head_lines_count: int
    tail_lines_count: int
    hidden_lines_count: int
    hidden_chars_count: int
    horiz_truncated_line_indices: List[int]

    @property
    def is_truncated(self) -> bool:
        return self.hidden_lines_count > 0 or bool(self.horiz_truncated_line_indices)

    def __iter__(self):
        # Support unpacking into (text, is_truncated)
        yield self.text
        yield self.is_truncated


LINE_SEP = "\n"


def split_text_lines(text: str) -> list[str]:
    # lines = text.split(LINE_SEP)
    lines = text.splitlines(keepends=True)
    return lines


def merge_text_lines(lines: list[str]) -> str:
    # text = LINE_SEP.join(lines)
    text = "".join(lines)
    return text


def truncate_text(
    orig_text: str, max_total_len: int, max_line_len: Optional[int] = None
) -> TruncationResult:
    # 1. Check total length first
    if len(orig_text) <= max_total_len:
        orig_lines = split_text_lines(orig_text)
        return TruncationResult(
            text=orig_text,
            orig_text=orig_text,
            head_lines_count=len(orig_lines),
            tail_lines_count=0,
            hidden_lines_count=0,
            hidden_chars_count=0,
            horiz_truncated_line_indices=[],
        )

    # 2. Horizontal Truncation (per line)
    orig_lines = split_text_lines(orig_text)
    horiz_processed_lines = []
    horiz_truncated_line_flags = []
    total_len_after_horiz = 0
    for line in orig_lines:
        if max_line_len is not None:
            line, truncated_flag = middle_truncate(line, max_line_len, False)
        else:
            truncated_flag = False
        horiz_processed_lines.append(line)
        horiz_truncated_line_flags.append(truncated_flag)
        total_len_after_horiz += len(line)
    horiz_processed_lines_count = len(horiz_processed_lines)

    # 3. Check if vertical truncation is needed after horizontal processing
    if total_len_after_horiz <= max_total_len:
        return TruncationResult(
            text=merge_text_lines(horiz_processed_lines),
            orig_text=orig_text,
            head_lines_count=horiz_processed_lines_count,
            tail_lines_count=0,
            hidden_lines_count=0,
            hidden_chars_count=0,
            horiz_truncated_line_indices=[
                i for i, t in enumerate(horiz_truncated_line_flags) if t
            ],
        )

    # 4. Calculate vertical cuts (Line-based to avoid cutting inside a line)
    head_threshold = max_total_len // 2
    tail_threshold = max_total_len - head_threshold

    # i_head: find the first line that ends at or exceeds head_threshold
    # We include this entire line in the head part.
    i_head = -1
    head_len = 0
    for i, line in enumerate(horiz_processed_lines):
        head_len += len(line)
        if head_len >= head_threshold:
            i_head = i
            break

    # i_tail: find the first line (from end) that makes the tail reach or exceed tail_threshold
    # We include this entire line in the tail part.
    i_tail = horiz_processed_lines_count
    tail_len = 0
    for i in range(horiz_processed_lines_count - 1, -1, -1):
        tail_len += len(horiz_processed_lines[i])
        if tail_len >= tail_threshold:
            i_tail = i
            break

    hidden_lines_count = i_tail - i_head - 1

    # 5. Check if we still need to truncate vertically
    # If head and tail sections overlap or touch, no lines are hidden.
    if hidden_lines_count <= 0:
        return TruncationResult(
            text=merge_text_lines(horiz_processed_lines),
            orig_text=orig_text,
            head_lines_count=horiz_processed_lines_count,
            tail_lines_count=0,
            hidden_lines_count=0,
            hidden_chars_count=0,
            horiz_truncated_line_indices=[
                i for i, t in enumerate(horiz_truncated_line_flags) if t
            ],
        )

    # 6. Construct Final Result (Line-based)
    head_part = horiz_processed_lines[: i_head + 1]
    tail_part = horiz_processed_lines[i_tail:]
    head_text = merge_text_lines(head_part)
    tail_text = merge_text_lines(tail_part)

    hidden_chars_count = total_len_after_horiz - head_len - tail_len
    marker_line = (
        f"...[TRUNCATED {hidden_lines_count} lines, {hidden_chars_count} chars]...\n"
    )
    final_text = head_text + marker_line + tail_text

    # 7. Map Truncated Indices
    truncated_line_indices = []

    # Head part indices (0 to i_head)
    for k in range(0, i_head + 1):
        if horiz_truncated_line_flags[k]:
            truncated_line_indices.append(k)

    # Tail part indices (i_tail to end)
    # We use negative indexing relative to the end of the final text.
    # -1 is the last line, -2 is the second to last, etc.
    for k in range(i_tail, horiz_processed_lines_count):
        if horiz_truncated_line_flags[k]:
            truncated_line_indices.append(k - horiz_processed_lines_count)

    return TruncationResult(
        text=final_text,
        orig_text=orig_text,
        head_lines_count=len(head_part),
        tail_lines_count=len(tail_part),
        hidden_lines_count=hidden_lines_count,
        hidden_chars_count=hidden_chars_count,
        horiz_truncated_line_indices=truncated_line_indices,
    )
