# Copyright (C) 2026 GoodData Corporation

"""Chunked stream reader with oversized-line handling."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from aida_mcp.utils.artifacts import get_artifacts_root


async def read_stream_chunked(
    stream: asyncio.StreamReader | None,
    *,
    stream_name: str,
    log_handle: Any,
    write_log_line: Callable[[Any, str, str], None],
    on_line: Callable[[str], Awaitable[None]] | None,
    max_capture_lines: int,
    max_capture_chars: int,
    max_line_bytes: int,
    sink: Any,
    chars_counter: int,
) -> int:
    """Read a stream safely, handling oversized lines and returning updated char count."""
    if not stream:
        return chars_counter

    buffer = bytearray()
    oversize = False
    oversize_handle: Any | None = None
    oversize_path: Path | None = None
    oversize_bytes = 0

    async def append_line(line: str, *, parse: bool) -> None:
        nonlocal chars_counter

        # Protocol line suppression:
        # `AIDA_EVENT ...` lines are used for streaming progress updates and should NOT
        # pollute logs/captured stdout/stderr. They still must be parsed for progress forwarding.
        is_progress_event = line.startswith("AIDA_EVENT ")

        if parse and on_line is not None:
            try:
                await on_line(line)
            except Exception:
                return

        if is_progress_event:
            return

        write_log_line(log_handle, stream_name, line)
        sink.append(line)
        chars_counter += len(line) + 1
        while len(sink) > max_capture_lines or chars_counter > max_capture_chars:
            removed = sink.popleft()
            chars_counter -= len(removed) + 1

    def start_oversize(chunk: bytes) -> None:
        nonlocal oversize, oversize_handle, oversize_path, oversize_bytes
        try:
            root = get_artifacts_root() / "oversized-lines"
            root.mkdir(parents=True, exist_ok=True)
            timestamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S-%f")
            oversize_path = root / f"{timestamp}-{stream_name}.log"
            oversize_handle = oversize_path.open("ab")
            oversize_handle.write(chunk)
        except OSError:
            oversize_handle = None
            oversize_path = None
        oversize_bytes = len(chunk)
        oversize = True

    async def finish_oversize() -> None:
        nonlocal oversize, oversize_handle, oversize_path, oversize_bytes
        if oversize_handle and oversize_path:
            oversize_handle.flush()
            oversize_handle.close()
            placeholder = f"[{stream_name}] Oversized line ({oversize_bytes} bytes) written to {oversize_path}"
            await append_line(placeholder, parse=False)
        elif oversize_bytes > 0:
            placeholder = f"[{stream_name}] Oversized line ({oversize_bytes} bytes) discarded"
            await append_line(placeholder, parse=False)
        oversize = False
        oversize_handle = None
        oversize_path = None
        oversize_bytes = 0

    while True:
        chunk = await stream.read(4096)
        if not chunk:
            break

        if oversize:
            nl_index = chunk.find(b"\n")
            if nl_index == -1:
                if oversize_handle:
                    oversize_handle.write(chunk)
                oversize_bytes += len(chunk)
                continue
            if oversize_handle:
                oversize_handle.write(chunk[: nl_index + 1])
            oversize_bytes += nl_index + 1
            await finish_oversize()
            buffer = bytearray(chunk[nl_index + 1 :])
        else:
            buffer.extend(chunk)

        while True:
            nl_index = buffer.find(b"\n")
            if nl_index != -1:
                line_bytes = buffer[: nl_index + 1]
                del buffer[: nl_index + 1]
                line = line_bytes.decode("utf-8", errors="replace").rstrip()
                await append_line(line, parse=True)
                continue

            if len(buffer) > max_line_bytes:
                start_oversize(bytes(buffer))
                buffer.clear()
            break

    if oversize:
        await finish_oversize()
    elif buffer:
        line = buffer.decode("utf-8", errors="replace").rstrip()
        await append_line(line, parse=True)

    return chars_counter
