"""JSON dataset implementation for semi-structured data storage.

Provides JSONDataset for loading and storing semi-structured JSON data.
Unlike SQLiteDataset which enforces a strict schema, JSONDataset works with
flexible JSON structures.
"""

import json
from typing import Any

from agentbyte.dataset.base import (
    BaseDataset,
    Column,
    DatasetConfig,
    Schema,
)


class JSONDataset(BaseDataset):
    """JSON-based dataset implementation for semi-structured data.
    
    Handles loading and storing JSON data files. Unlike SQLiteDataset,
    JSONDataset does not enforce a strict schema and allows flexible
    JSON structures.
    
    Attributes:
        dataset_id: Identifier (e.g., "documents/wiki")
        config: Dataset configuration
        data_text: Raw data.json content
    """

    def __init__(
        self,
        dataset_id: str,
        config: DatasetConfig,
        data_text: str,
        source_uri: str,
    ) -> None:
        """Initialize JSON dataset.
        
        Args:
            dataset_id: Dataset identifier
            config: DatasetConfig object
            data_text: Raw ``data.json`` content
            source_uri: Human-readable dataset source location
        """
        self.dataset_id = dataset_id
        self.config = config
        self.data_text = data_text
        self.source_uri = source_uri
        self._data: dict[str, Any] | None = None
        self._schema: Schema | None = None

    @property
    def engine(self) -> str:
        """Get the database engine name."""
        return self.config.engine

    @property
    def table_name(self) -> str:
        """Get the table name (used as dataset identifier in JSON)."""
        return self.config.table_name

    @property
    def schema(self) -> Schema | None:
        """Get the schema (inferred from data on first connect() call).
        
        Returns a generic Schema object with inferred column definitions.
        """
        return self._schema

    def connect(self) -> dict[str, Any]:
        """Get the loaded JSON data as a dictionary.
        
        Loads data from data.json on first call and infers schema from records.
        
        Returns:
            dict with structure: {"records": [...], "metadata": {...}}
        """
        if self._data is not None:
            return self._data

        if self.config.engine != "json":
            raise NotImplementedError(
                f"Engine '{self.config.engine}' not supported by JSONDataset. "
                "Currently 'json' is the only supported engine."
            )

        data = json.loads(self.data_text)

        # Normalize data structure
        # Support both {"records": [...]} and raw array []
        if isinstance(data, list):
            data = {"records": data}
        elif not isinstance(data, dict):
            raise ValueError(
                "Data file must contain either a list or a dict with 'records' key "
                f"for dataset: {self.source_uri}/data.json"
            )

        # Infer schema from first record (if records exist)
        records = data.get("records", [])
        if records and isinstance(records[0], dict):
            columns = [
                Column(name=key, type="string")  # JSON treats all values as JSON types
                for key in records[0].keys()
            ]
            self._schema = Schema(columns=columns)

        self._data = data
        return data

    def __repr__(self) -> str:
        """String representation of dataset."""
        return (
            f"JSONDataset(id='{self.dataset_id}', engine='{self.engine}', "
            f"table='{self.table_name}')"
        )
