import json
import logging
from typing import Iterator, Dict, Any, Union, List, TextIO, Iterable
from pathlib import Path

from atomingest.normalization.steps import BaseNormalizerStep
from atomingest.normalization.sources.base import DataSource

logger = logging.getLogger(__name__)


class JSONDataSource(DataSource):
    """
    Concrete Strategy for JSON files
    """

    def read(
        self,
        source: Union[str, Path, TextIO, Iterable[str]],
        line_steps: List[BaseNormalizerStep] = None,
        report=None,
    ) -> Iterator[Dict[str, Any]]:
        if line_steps:
            logger.warning(
                "Line steps are not supported for JSONDataSource and will be ignored."
            )

        data = None
        if isinstance(source, (str, Path)):
            with open(source, "r", encoding="utf-8") as f:
                data = json.load(f)
        elif hasattr(source, "read"):
            data = json.load(source)  # type: ignore[arg-type]
        else:
            raise ValueError(
                "Invalid Source type for JSONDataSource. JSON source must be a file path or file-like object."
            )

        if isinstance(data, dict):
            if report:
                report.increment_total_lines()
            yield data
        elif isinstance(data, list):
            for item in data:
                if report:
                    report.increment_total_lines()
                yield item
        else:
            raise ValueError("JSON source must contain a list or a dict")
