"""
Simple streaming JSON array parser for JAF.

This module provides a lightweight streaming parser that extracts JSON objects
from a JSON array without loading the entire array into memory. It works by
counting braces to identify complete objects.
"""

import json
import gzip
from typing import Generator, Dict, Any, IO, Union, Optional
import io


def stream_json(
    file_obj: Union[IO[str], IO[bytes]],
    chunk_size: int = 65536,
    initial_object_size: int = 4096,
) -> Generator[Dict[str, Any], None, None]:
    """
    Streaming JSON parser that handles both arrays and single objects.

    This parser can handle:
    - JSON array of objects: [{...}, {...}, ...]
    - Single JSON object: {...}

    For arrays, it streams each object. For single objects, it yields the object once.

    Args:
        file_obj: File object (text or binary mode)
        chunk_size: Size of chunks to read from file (default 64KB)
        initial_object_size: Initial size for object buffer (default 4KB)

    Yields:
        Dict[str, Any]: Each JSON object

    Raises:
        ValueError: If the file contains invalid JSON
    """
    # Handle both text and binary files
    if hasattr(file_obj, "mode") and "b" in file_obj.mode:
        # Binary mode
        read_chunk = lambda: file_obj.read(chunk_size).decode("utf-8", errors="replace")
    else:
        # Text mode
        read_chunk = lambda: file_obj.read(chunk_size)

    # Read first chunk
    chunk = read_chunk()
    if not chunk:
        return

    # Find first non-whitespace character
    pos = skip_whitespace(chunk, 0)
    if pos >= len(chunk):
        return

    # Check if it's an array
    if chunk[pos] != "[":
        # Not an array - read entire value and yield it
        all_chunks = [chunk]
        while True:
            next_chunk = read_chunk()
            if not next_chunk:
                break
            all_chunks.append(next_chunk)

        full_content = "".join(all_chunks)
        try:
            value = json.loads(full_content)
            yield value
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON: {e}")
        return

    # It's an array - stream each value
    pos += 1  # Skip the '['

    # Create chunk iterator
    def chunk_iterator():
        nonlocal chunk, pos
        while True:
            next_chunk = read_chunk()
            if not next_chunk:
                break
            yield next_chunk

    chunks = chunk_iterator()

    while True:
        # Skip whitespace and commas
        while pos < len(chunk) and chunk[pos] in " \t\n\r,":
            pos += 1

        # Need more data?
        if pos >= len(chunk):
            try:
                chunk = next(chunks)
                pos = 0
                continue
            except StopIteration:
                return

        # Check for end of array
        if chunk[pos] == "]":
            return

        # Parse next value based on first character
        first_char = chunk[pos]

        try:
            if first_char == "{":
                # Object
                value_str, chunk, pos = extract_balanced_value(
                    chunks, chunk, pos, "{", "}"
                )
            elif first_char == "[":
                # Nested array
                value_str, chunk, pos = extract_balanced_value(
                    chunks, chunk, pos, "[", "]"
                )
            elif first_char == '"':
                # String
                value_str, chunk, pos = extract_string_value(chunks, chunk, pos)
            else:
                # Number, boolean, or null
                value_str, chunk, pos = extract_literal_value(chunks, chunk, pos)

            # Parse and yield the value
            value = json.loads(value_str)
            yield value

        except (json.JSONDecodeError, ValueError):
            # Skip malformed values, try to recover
            # Find next comma or closing bracket
            while pos < len(chunk) and chunk[pos] not in ",]":
                pos += 1


def skip_whitespace(text: str, start: int) -> int:
    """
    Skip whitespace characters and return position of next non-whitespace.

    Args:
        text: String to scan
        start: Starting position

    Returns:
        Position of next non-whitespace character, or len(text) if all whitespace
    """
    # Use built-in string methods for efficiency
    remainder = text[start:]
    stripped = remainder.lstrip()
    if stripped:
        return start + (len(remainder) - len(stripped))
    return len(text)


def extract_balanced_value(
    chunks_iter, initial_chunk, start_pos, open_char, close_char
):
    """
    Extract a balanced JSON value (object or array) by counting brackets/braces.

    Returns: (value_str, remaining_chunk, next_pos)
    """
    parts = []
    chunk = initial_chunk
    pos = start_pos

    in_string = False
    escape_next = False
    depth = 0

    while True:
        while pos < len(chunk):
            char = chunk[pos]

            # Handle string state
            if not escape_next:
                if char == '"' and not in_string:
                    in_string = True
                elif char == '"' and in_string:
                    in_string = False
                elif char == "\\" and in_string:
                    escape_next = True
            else:
                escape_next = False

            # Count brackets/braces when not in string
            if not in_string:
                if char == open_char:
                    depth += 1
                elif char == close_char:
                    depth -= 1
                    if depth == 0:
                        # Complete value found
                        parts.append(chunk[start_pos : pos + 1])
                        return "".join(parts), chunk, pos + 1

            pos += 1

        # Need more data
        parts.append(chunk[start_pos:])
        try:
            chunk = next(chunks_iter)
            start_pos = 0
            pos = 0
        except StopIteration:
            raise ValueError(
                f"Unexpected end of input while parsing {open_char}...{close_char}"
            )


def extract_string_value(chunks_iter, initial_chunk, start_pos):
    """
    Extract a JSON string value.

    Returns: (value_str, remaining_chunk, next_pos)
    """
    parts = []
    chunk = initial_chunk
    pos = start_pos + 1  # Skip opening quote

    escape_next = False

    while True:
        while pos < len(chunk):
            char = chunk[pos]

            if not escape_next:
                if char == '"':
                    # String complete
                    parts.append(chunk[start_pos : pos + 1])
                    return "".join(parts), chunk, pos + 1
                elif char == "\\":
                    escape_next = True
            else:
                escape_next = False

            pos += 1

        # Need more data
        parts.append(chunk[start_pos:])
        try:
            chunk = next(chunks_iter)
            start_pos = 0
            pos = 0
        except StopIteration:
            raise ValueError("Unexpected end of input while parsing string")


def extract_literal_value(chunks_iter, initial_chunk, start_pos):
    """
    Extract a JSON literal (number, true, false, null).

    Returns: (value_str, remaining_chunk, next_pos)
    """
    chunk = initial_chunk
    pos = start_pos

    # Find the end of the literal (whitespace, comma, or closing bracket)
    while pos < len(chunk):
        char = chunk[pos]
        if char in " \t\n\r,]}":
            return chunk[start_pos:pos], chunk, pos
        pos += 1

    # Literal continues into next chunk
    parts = [chunk[start_pos:]]

    try:
        while True:
            chunk = next(chunks_iter)
            pos = 0
            while pos < len(chunk):
                char = chunk[pos]
                if char in " \t\n\r,]}":
                    parts.append(chunk[:pos])
                    return "".join(parts), chunk, pos
                pos += 1
            parts.append(chunk)
    except StopIteration:
        # End of input
        return "".join(parts), "", 0


def stream_jsonl(file_obj: Union[IO[str], IO[bytes]]) -> Generator[Any, None, None]:
    """
    Stream JSON values from a JSONL (JSON Lines) file.

    Args:
        file_obj: An open file object containing JSONL data

    Yields:
        Any: JSON values (one per line)
    """
    for line in file_obj:
        if isinstance(line, bytes):
            line = line.decode("utf-8", errors="replace")
        line = line.strip()
        if line:
            try:
                value = json.loads(line)
                yield value
            except json.JSONDecodeError:
                # Skip invalid lines
                pass


def stream_json_file(
    file_path: str, format: Optional[str] = None
) -> Generator[Dict[str, Any], None, None]:
    """
    Stream JSON objects from a file, automatically handling compression and format.

    Args:
        file_path: Path to the file
        format: Force a specific format ('json', 'jsonl', or None for auto-detect)

    Yields:
        Dict[str, Any]: Each JSON object in the file
    """
    # Determine format from extension if not specified
    if format is None:
        if file_path.endswith(".jsonl") or file_path.endswith(".jsonl.gz"):
            format = "jsonl"
        else:
            format = "json"

    # Handle gzip compression
    if file_path.endswith(".gz"):
        file_obj = gzip.open(file_path, "rt", encoding="utf-8")
    else:
        file_obj = open(file_path, "r", encoding="utf-8")

    try:
        if format == "jsonl":
            yield from stream_jsonl(file_obj)
        else:
            yield from stream_json_array(file_obj)
    finally:
        file_obj.close()


def stream_json_collection(
    source: Dict[str, Any],
) -> Generator[Dict[str, Any], None, None]:
    """
    Stream JSON objects from various collection sources.

    Args:
        source: Collection source descriptor with 'type' and type-specific fields

    Yields:
        Dict[str, Any]: Each JSON object in the collection
    """
    source_type = source.get("type") or source.get("source_type")

    if source_type == "jsonl":
        path = source.get("path")
        if path:
            yield from stream_json_file(path, format="jsonl")

    elif source_type == "json_array":
        path = source.get("path")
        if path:
            yield from stream_json_file(path, format="json")

    elif source_type == "directory":
        files = source.get("files", [])
        for file_path in files:
            if file_path.endswith(".jsonl") or file_path.endswith(".jsonl.gz"):
                yield from stream_json_file(file_path, format="jsonl")
            else:
                yield from stream_json_file(file_path, format="json")

    elif source_type == "buffered_stdin" or source_type == "in_memory":
        # These are already in memory, just yield them
        content = source.get("content") or source.get("data", [])
        for item in content:
            if isinstance(item, dict):
                yield item

    else:
        raise ValueError(f"Unknown source type: {source_type}")
