"""
Custom JSON document loader.

This loader reads JSON files and extracts all string values for text search.
It doesn't require the jq package like LangChain's JSONLoader.
"""

import json
from pathlib import Path
from typing import Any, Iterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document
from loguru import logger


def extract_strings_from_json(
    data: Any, path: str = ""
) -> list[tuple[str, str]]:
    """
    Recursively extract all string values from JSON data.

    Returns list of (path, value) tuples where path shows the JSON path.
    """
    results = []

    if isinstance(data, str):
        results.append((path, data))
    elif isinstance(data, dict):
        for key, value in data.items():
            new_path = f"{path}.{key}" if path else key
            results.extend(extract_strings_from_json(value, new_path))
    elif isinstance(data, list):
        for i, item in enumerate(data):
            new_path = f"{path}[{i}]"
            results.extend(extract_strings_from_json(item, new_path))
    # Skip numbers, booleans, None - they're not useful for text search

    return results


class SimpleJSONLoader(BaseLoader):
    """
    Load JSON files and extract all string values for text search.

    This loader recursively extracts all string values from the JSON structure,
    making them searchable. Numbers and booleans are skipped.
    """

    def __init__(
        self,
        file_path: str | Path,
        encoding: str = "utf-8",
        include_paths: bool = True,
    ):
        """
        Initialize the JSON loader.

        Args:
            file_path: Path to the JSON file
            encoding: File encoding (default: utf-8)
            include_paths: If True, include JSON paths with values (default: True)
        """
        self.file_path = Path(file_path)
        self.encoding = encoding
        self.include_paths = include_paths

    def lazy_load(self) -> Iterator[Document]:
        """
        Lazily load documents from the JSON file.

        Yields:
            Document objects containing the extracted text
        """
        try:
            with open(self.file_path, encoding=self.encoding) as f:
                content = f.read()

            try:
                data = json.loads(content)
            except json.JSONDecodeError:
                logger.exception(f"Invalid JSON in {self.file_path}")
                # Still yield the raw content if parsing fails
                yield Document(
                    page_content=content,
                    metadata={
                        "source": str(self.file_path),
                        "parse_error": True,
                    },
                )
                return

            # Extract all string values
            string_pairs = extract_strings_from_json(data)

            if self.include_paths:
                # Format as "path: value" for each string
                lines = [
                    f"{path}: {value}" for path, value in string_pairs if value
                ]
            else:
                # Just the values
                lines = [value for _, value in string_pairs if value]

            text = "\n".join(lines)

            yield Document(
                page_content=text,
                metadata={
                    "source": str(self.file_path),
                    "file_type": "json",
                    "string_count": len(string_pairs),
                },
            )

        except Exception:
            logger.exception(f"Error loading JSON file: {self.file_path}")
            raise

    def load(self) -> list[Document]:
        """Load all documents from the JSON file."""
        return list(self.lazy_load())


def extract_text_from_json(content: bytes, encoding: str = "utf-8") -> str:
    """
    Extract text from JSON bytes content.

    This is a convenience function for use with bytes (e.g., from file uploads).

    Args:
        content: JSON content as bytes
        encoding: Text encoding (default: utf-8)

    Returns:
        Extracted text with JSON paths and values
    """
    try:
        text = content.decode(encoding)
    except UnicodeDecodeError:
        text = content.decode(encoding, errors="ignore")
        logger.warning(
            "JSON content had encoding errors, some characters may be lost"
        )

    try:
        data = json.loads(text)
        string_pairs = extract_strings_from_json(data)
        lines = [f"{path}: {value}" for path, value in string_pairs if value]
        return "\n".join(lines)
    except json.JSONDecodeError:
        logger.exception("Invalid JSON content")
        return text
