"""
Generic JSON stream accumulator for line-oriented process output.
Task-agnostic: decodes either brace-balanced JSON objects spanning lines
or newline-delimited JSON (NDJSON).

Usage:
    acc = JsonAccumulator(ndjson=False)
    for line in lines:
        for obj in acc.feed(line):
            handle(obj)
"""

from __future__ import annotations

import json
from typing import Optional
from hil_testbench.data_structs.parameters import ParametersSchema


class JsonAccumulator:
    """Accumulates streaming JSON objects, validates against schema if provided."""

    def __init__(
        self, ndjson: bool = False, schema: Optional[ParametersSchema] = None
    ) -> None:
        self.ndjson = ndjson
        self.schema = schema
        self._buffer: list[str] = []
        self._brace_balance = 0

    def reset(self) -> None:
        """Reset the internal buffer and brace counter."""
        self._buffer.clear()
        self._brace_balance = 0

    def feed(self, line: str) -> list[dict[str, object]]:
        """Feed one line of text; return zero or more parsed JSON objects.

        This method is resilient to malformed input: if decoding fails,
        the offending chunk is discarded and parsing continues.
        """
        out: list[dict[str, object]] = []
        if not line:
            return out

        if self.ndjson:
            line = line.strip()
            if not line:
                return out
            try:
                obj = json.loads(line)
                if isinstance(obj, dict):
                    out.append(obj)
            except json.JSONDecodeError:
                pass
            return out

        # Brace-balanced accumulation for multi-line objects
        open_count = line.count("{")
        close_count = line.count("}")
        if open_count:
            self._brace_balance += open_count
        self._buffer.append(line.rstrip("\n"))
        if close_count:
            self._brace_balance -= close_count

        if self._brace_balance != 0:
            return out  # object not complete yet

        # We have a complete object in the buffer
        raw = "\n".join(self._buffer)
        self.reset()
        try:
            obj = json.loads(raw)
            if isinstance(obj, dict):
                out.append(obj)
        except json.JSONDecodeError:
            # discard malformed chunk
            pass
        return out

    def _parse_json(self, line: str) -> list[dict[str, object]]:
        """Parse the JSON or NDJSON line, returning a list of objects."""
        out: list[dict[str, object]] = []
        if not line:
            return out

        if self.ndjson:
            line = line.strip()
            if not line:
                return out
            try:
                obj = json.loads(line)
                if isinstance(obj, dict):
                    out.append(obj)
            except json.JSONDecodeError:
                pass
            return out

        # Brace-balanced accumulation for multi-line objects
        open_count = line.count("{")
        close_count = line.count("}")
        if open_count:
            self._brace_balance += open_count
        self._buffer.append(line.rstrip("\n"))
        if close_count:
            self._brace_balance -= close_count

        if self._brace_balance != 0:
            return out  # object not complete yet

        # We have a complete object in the buffer
        raw = "\n".join(self._buffer)
        self.reset()
        try:
            obj = json.loads(raw)
            if isinstance(obj, dict):
                out.append(obj)
        except json.JSONDecodeError:
            # discard malformed chunk
            pass
        return out

    def feed_with_validation(self, line: str) -> list[dict[str, object]]:
        """Feed a line of JSON or NDJSON, return parsed objects. Validate if schema is set."""
        objs = self._parse_json(line)
        if self.schema:
            for obj in objs:
                for field in obj.keys():
                    if field not in [f.name for f in self.schema.fields]:
                        raise ValueError(f"Unknown field '{field}' not in schema.")
        return objs
