"""JSON formatter for dlist"""

import json
from .base import BaseFormatter


class JSONFormatter(BaseFormatter):
    """JSON formatter for Dlist objects

    Supports flat and categorized output.

    Examples::

        >>> fmt = JSONFormatter()
        >>> print(fmt.format(d, keys=['name', 'age']))
        >>> fmt.save(d, 'output.json')
        >>> print(fmt.format_categorized(d, ['type'], ['name']))
    """

    def format(self, dlist, keys=[], **kwargs):
        """Format Dlist as JSON

        Parameters:
            dlist: Dlist object to format
            keys (list): If provided, only include these keys in output
            **kwargs: JSON options (indent, sort_keys)

        Returns:
            str: JSON formatted string
        """
        indent = kwargs.get('indent', 2)
        sort_keys = kwargs.get('sort_keys', False)

        if keys:
            filtered_data = []
            for item in dlist.l:
                filtered_item = {}
                for key in keys:
                    if key in item:
                        filtered_item[key] = item[key]
                filtered_data.append(filtered_item)
            return json.dumps(filtered_data, indent=indent, sort_keys=sort_keys)
        else:
            return json.dumps(dlist.l, indent=indent, sort_keys=sort_keys)

    def format_categorized(self, dlist, ctree, keys, titles={}, **kwargs):
        """Format with categories as nested JSON

        Parameters:
            dlist: Dlist object to format
            ctree (list): List of keys used as categories
            keys (list): List of keys to include
            titles (dict): Not used for JSON
            **kwargs: JSON options

        Returns:
            str: Categorized JSON output
        """
        indent = kwargs.get('indent', 2)

        if not ctree:
            return self.format(dlist, keys=keys, **kwargs)

        # Create nested structure based on categories
        partition = dlist.partition(ctree[0])

        if len(ctree) == 1:
            # Last level - return the data
            result = {}
            for cat_val, cat_dlist in partition.items():
                if keys:
                    filtered = []
                    for item in cat_dlist.l:
                        filtered_item = {k: item.get(k) for k in keys if k in item}
                        filtered.append(filtered_item)
                    result[str(cat_val)] = filtered
                else:
                    result[str(cat_val)] = cat_dlist.l
            return json.dumps(result, indent=indent)
        else:
            # Recursive categorization
            result = {}
            for cat_val, cat_dlist in partition.items():
                result[str(cat_val)] = json.loads(
                    self.format_categorized(cat_dlist, ctree[1:], keys, titles, **kwargs)
                )
            return json.dumps(result, indent=indent)

    def save(self, dlist, filepath, **kwargs):
        """Save Dlist to a JSON file

        Parameters:
            dlist: Dlist object to format
            filepath (str): Output file path
            **kwargs: Passed to format() (keys, indent, sort_keys)
        """
        content = self.format(dlist, **kwargs)
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)
