"""
Streaming collection loader for JAF.

This module provides a composable, streaming-only loader system where each
loader is a generator that transforms its input stream. Loaders can be
chained together to form processing pipelines.

Example pipeline:
    file → gzip → jsonl → objects
    http → tar → zip → csv → objects
"""

import json
import gzip
import csv
import tarfile
import zipfile
from pathlib import Path
from typing import Generator, Dict, Any, Union, IO, Iterator
import io


# Type aliases
JsonObject = Dict[str, Any]
ByteStream = Generator[bytes, None, None]
TextStream = Generator[str, None, None]
ObjectStream = Generator[JsonObject, None, None]


class StreamingLoader:
    """
    Registry and dispatcher for streaming loaders.

    Each loader is a generator function that takes a source descriptor
    and yields items. Loaders can reference other loaders via inner_source
    to create processing pipelines.
    """

    def __init__(self):
        self._loaders = {}
        self._register_default_loaders()

    def register(self, name: str, loader_func):
        """Register a streaming loader function."""
        self._loaders[name] = loader_func

    def stream(self, source: Dict[str, Any]) -> ObjectStream:
        """
        Stream objects from a source descriptor.

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

        Yields:
            JSON objects from the source
        """
        source_type = source.get("type") or source.get("source_type")
        if not source_type:
            raise ValueError("Source descriptor missing 'type' field")

        loader = self._loaders.get(source_type)
        if not loader:
            raise ValueError(f"Unknown source type: {source_type}")

        yield from loader(self, source)

    def _register_default_loaders(self):
        """Register built-in loaders."""
        # File system
        self.register("file", stream_file)
        self.register("directory", stream_directory)

        # Decompression
        self.register("gzip", stream_gzip)
        self.register("tar", stream_tar)
        self.register("zip", stream_zip)

        # Parsers
        self.register("jsonl", parse_jsonl)
        self.register("json_array", parse_json_array)
        self.register("json_value", parse_json_value)
        self.register("csv", parse_csv)

        # Special
        self.register("memory", stream_memory)

        # Register codata loaders
        from .codata_loaders import register_codata_loaders

        register_codata_loaders(self)

        # Register lazy operation loaders
        from .lazy_ops_loader import register_lazy_ops_loaders

        register_lazy_ops_loaders(self)


# --- File System Loaders ---


def stream_file(loader: StreamingLoader, source: Dict[str, Any]) -> ByteStream:
    """Stream bytes from a file."""
    path = source.get("path")
    if not path:
        raise ValueError("File source missing 'path'")

    chunk_size = source.get("chunk_size", 65536)

    with open(path, "rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            yield chunk


def stream_directory(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """
    Stream objects from multiple files in a directory.

    Args:
        source: Dict with:
            - path: Directory path
            - pattern: Glob pattern (default: '*.json*')
            - recursive: Whether to search recursively (default: False)
            - file_as_object: If True, treat each JSON file as a single object (default: False)
    """
    path = source.get("path")
    if not path:
        raise ValueError("Directory source missing 'path'")

    pattern = source.get("pattern", "*.json*")
    recursive = source.get("recursive", False)
    file_as_object = source.get("file_as_object", False)

    from pathlib import Path

    base_path = Path(path)

    if recursive:
        files = base_path.rglob(pattern)
    else:
        files = base_path.glob(pattern)

    for file_path in files:
        # Build clean pipeline for each file
        file_source = {"type": "file", "path": str(file_path)}

        # Add decompression if needed
        if str(file_path).endswith(".gz"):
            file_source = {"type": "gzip", "inner_source": file_source}

        # Add parser based on extension and settings
        if ".jsonl" in str(file_path):
            file_source = {"type": "jsonl", "inner_source": file_source}
        elif file_as_object:
            # Treat entire file as one value
            file_source = {"type": "json_value", "inner_source": file_source}
        else:
            # Try to stream array elements (falls back to single value if not array)
            file_source = {"type": "json_array", "inner_source": file_source}

        yield from loader.stream(file_source)


# --- Decompression Loaders ---


def stream_gzip(loader: StreamingLoader, source: Dict[str, Any]) -> ByteStream:
    """Stream decompressed bytes from a gzipped source."""
    inner_source = source.get("inner_source")
    if not inner_source:
        raise ValueError("Gzip source missing 'inner_source'")

    # Get byte stream from inner source
    byte_stream = loader.stream(inner_source)

    # Create a file-like object from the byte stream
    class ByteStreamWrapper:
        def __init__(self, stream):
            self.stream = stream
            self.buffer = b""

        def read(self, size=-1):
            if size == -1:
                # Read all
                result = self.buffer
                for chunk in self.stream:
                    result += chunk
                self.buffer = b""
                return result

            # Read up to size bytes
            while len(self.buffer) < size:
                try:
                    chunk = next(self.stream)
                    self.buffer += chunk
                except StopIteration:
                    break

            result = self.buffer[:size]
            self.buffer = self.buffer[size:]
            return result

    wrapper = ByteStreamWrapper(byte_stream)
    with gzip.GzipFile(fileobj=wrapper, mode="rb") as gz:
        while True:
            chunk = gz.read(8192)
            if not chunk:
                break
            yield chunk


# --- Parser Loaders ---


def parse_jsonl(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Parse JSONL from a byte/text stream."""
    inner_source = source.get("inner_source")
    if not inner_source:
        raise ValueError("JSONL parser missing 'inner_source'")

    # Get stream from inner source
    stream = loader.stream(inner_source)

    # Build lines from chunks
    buffer = ""
    for chunk in stream:
        if isinstance(chunk, bytes):
            chunk = chunk.decode("utf-8", errors="replace")

        buffer += chunk
        lines = buffer.split("\n")

        # Process all complete lines
        for line in lines[:-1]:
            line = line.strip()
            if line:
                try:
                    value = json.loads(line)
                    yield value
                except json.JSONDecodeError:
                    pass  # Skip invalid lines

        # Keep the incomplete line in buffer
        buffer = lines[-1]

    # Process final line if any
    if buffer.strip():
        try:
            value = json.loads(buffer)
            yield value
        except json.JSONDecodeError:
            pass


def parse_json_array(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Parse JSON array from a byte/text stream using custom streaming parser."""
    inner_source = source.get("inner_source")
    if not inner_source:
        raise ValueError("JSON array parser missing 'inner_source'")

    # Get stream from inner source
    stream = loader.stream(inner_source)

    # Use our optimized JSON array streamer
    class StreamWrapper:
        def __init__(self, stream):
            self.stream = stream
            self.buffer = ""
            self.done = False

        def read(self, size=-1):
            if self.done:
                return ""

            if size == -1:
                # Read all
                result = self.buffer
                for chunk in self.stream:
                    if isinstance(chunk, bytes):
                        chunk = chunk.decode("utf-8", errors="replace")
                    result += chunk
                self.done = True
                return result

            # Read up to size characters
            while len(self.buffer) < size:
                try:
                    chunk = next(self.stream)
                    if isinstance(chunk, bytes):
                        chunk = chunk.decode("utf-8", errors="replace")
                    self.buffer += chunk
                except StopIteration:
                    self.done = True
                    break

            result = self.buffer[:size]
            self.buffer = self.buffer[size:]
            return result

    wrapper = StreamWrapper(stream)
    from .json_stream import stream_json

    yield from stream_json(wrapper)


def parse_json_value(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Parse JSON file as a single value (treat whole file as one collection item)."""
    inner_source = source.get("inner_source")
    if not inner_source:
        raise ValueError("JSON value parser missing 'inner_source'")

    # Get stream from inner source
    stream = loader.stream(inner_source)

    # Read entire content
    chunks = []
    for chunk in stream:
        if isinstance(chunk, bytes):
            chunk = chunk.decode("utf-8", errors="replace")
        chunks.append(chunk)

    content = "".join(chunks).strip()
    if content:
        try:
            value = json.loads(content)
            yield value
        except json.JSONDecodeError:
            pass  # Skip invalid JSON


def parse_csv(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Parse CSV into JSON objects."""
    inner_source = source.get("inner_source")
    if not inner_source:
        raise ValueError("CSV parser missing 'inner_source'")

    # CSV options
    delimiter = source.get("delimiter", ",")
    has_header = source.get("has_header", True)

    # Get stream from inner source
    stream = loader.stream(inner_source)

    # Convert byte stream to line stream
    buffer = ""
    lines = []

    for chunk in stream:
        if isinstance(chunk, bytes):
            chunk = chunk.decode("utf-8", errors="replace")

        buffer += chunk
        new_lines = buffer.split("\n")

        # Keep incomplete line in buffer
        lines.extend(new_lines[:-1])
        buffer = new_lines[-1]

    # Add final line if any
    if buffer:
        lines.append(buffer)

    # Parse CSV
    reader = csv.reader(lines, delimiter=delimiter)

    if has_header:
        try:
            headers = next(reader)
        except StopIteration:
            return

        for row in reader:
            if row:  # Skip empty rows
                # Convert to dict using headers
                obj = {}
                for i, value in enumerate(row):
                    if i < len(headers):
                        # Try to parse numbers
                        try:
                            if "." in value:
                                obj[headers[i]] = float(value)
                            else:
                                obj[headers[i]] = int(value)
                        except ValueError:
                            obj[headers[i]] = value
                yield obj
    else:
        # No header - use numeric keys
        for row in reader:
            if row:
                obj = {str(i): value for i, value in enumerate(row)}
                yield obj


def stream_memory(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Stream values from memory (for in-memory collections)."""
    data = source.get("data") or source.get("content", [])
    for item in data:
        yield item


# --- Archive Loaders (stubs for now) ---


def stream_tar(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Stream objects from tar archive members."""
    # TODO: Implement streaming tar extraction
    raise NotImplementedError("TAR streaming not yet implemented")


def stream_zip(loader: StreamingLoader, source: Dict[str, Any]) -> ObjectStream:
    """Stream objects from zip archive members."""
    # TODO: Implement streaming zip extraction
    raise NotImplementedError("ZIP streaming not yet implemented")
