"""JSON and JSONL data loaders.

Provides concrete implementations of DataLoader for JSON and JSONL file formats.
"""

import json

from agentbyte.dataset.base import DataLoader, DataRecord


class JSONLoader(DataLoader[DataRecord]):
    """Loader for JSON files containing arrays of records.
    
    Expects a JSON file with an array of objects, where each object
    can optionally have 'id', 'data', and 'metadata' fields.
    
    Example JSON structure:
        [
            {"id": "doc1", "data": {"content": "..."}, "metadata": {"source": "file"}},
            {"id": "doc2", "data": {"content": "..."}, "metadata": {"source": "file"}},
        ]
    """

    async def load(self) -> list[DataRecord]:
        """Load records from a JSON file.
        
        Returns:
            List of DataRecord objects
            
        Raises:
            FileNotFoundError: If JSON file doesn't exist
            json.JSONDecodeError: If file is not valid JSON
            ValueError: If records don't have required fields
        """
        if not self.source.exists():
            raise FileNotFoundError(f"JSON file not found: {self.source}")

        with open(self.source, "r") as f:
            data = json.load(f)

        if not isinstance(data, list):
            raise ValueError("JSON file must contain an array of records")

        records = []
        for idx, item in enumerate(data):
            if not isinstance(item, dict):
                raise ValueError(f"Record {idx} is not a dictionary")

            # Use provided id or generate one
            record_id = item.get("id", f"record_{idx}")
            record_data = item.get("data", item)  # Use whole item if no 'data' field
            record_metadata = item.get("metadata", {})

            records.append(
                DataRecord(
                    id=record_id,
                    data=record_data,
                    metadata=record_metadata,
                )
            )

        return records


class JSONLLoader(DataLoader[DataRecord]):
    """Loader for JSONL (JSON Lines) files.
    
    Expects a file where each line is a valid JSON object (no commas between lines).
    Each line becomes one record.
    
    Example JSONL structure:
        {"id": "doc1", "data": {"content": "..."}, "metadata": {"source": "file"}}
        {"id": "doc2", "data": {"content": "..."}, "metadata": {"source": "file"}}
    """

    async def load(self) -> list[DataRecord]:
        """Load records from a JSONL file.
        
        Returns:
            List of DataRecord objects
            
        Raises:
            FileNotFoundError: If JSONL file doesn't exist
            json.JSONDecodeError: If a line is not valid JSON
            ValueError: If records don't have required fields
        """
        if not self.source.exists():
            raise FileNotFoundError(f"JSONL file not found: {self.source}")

        records = []
        with open(self.source, "r") as f:
            for line_num, line in enumerate(f, start=1):
                line = line.strip()
                if not line:
                    continue  # Skip empty lines

                try:
                    item = json.loads(line)
                except json.JSONDecodeError as e:
                    raise ValueError(f"Invalid JSON on line {line_num}: {e}")

                if not isinstance(item, dict):
                    raise ValueError(f"Line {line_num} is not a JSON object")

                # Use provided id or generate one
                record_id = item.get("id", f"record_{line_num}")
                record_data = item.get("data", item)
                record_metadata = item.get("metadata", {})

                records.append(
                    DataRecord(
                        id=record_id,
                        data=record_data,
                        metadata=record_metadata,
                    )
                )

        return records


class DirectoryLoader(DataLoader[DataRecord]):
    """Loader for directories containing JSON/JSONL files.
    
    Scans a directory for all .json and .jsonl files and loads them.
    Each file is loaded as a collection of records.
    """

    async def load(self) -> list[DataRecord]:
        """Load all JSON/JSONL files from a directory.
        
        Returns:
            Combined list of DataRecord objects from all files
            
        Raises:
            FileNotFoundError: If directory doesn't exist
            ValueError: If files contain invalid data
        """
        if not self.source.is_dir():
            raise FileNotFoundError(f"Directory not found: {self.source}")

        all_records: list[DataRecord] = []

        # Load all JSON files
        for json_file in self.source.glob("*.json"):
            loader = JSONLoader(json_file)
            records = await loader.load()
            all_records.extend(records)

        # Load all JSONL files
        for jsonl_file in self.source.glob("*.jsonl"):
            loader = JSONLLoader(jsonl_file)
            records = await loader.load()
            all_records.extend(records)

        return all_records
