from __future__ import annotations

from typing import Any, List, Optional

from jsonpath_ng import parse
from jsonpath_ng.ext import parse as ext_parse
from jsonpath_ng.exceptions import JsonPathLexerError, JsonPathParserError

from app.core.logger import get_logger

logger = get_logger(__name__)


class KeysExtractor:
    """Nested JSON extractor using jsonpath-ng.

    Supports two modes:
      1. key_names — recursive lookup by key name at any depth
      2. query — extended JSONPath expression with full path output
    """

    def __init__(
        self,
        data: dict | list,
        key_names: Optional[List[str]] = None,
        result_limit: Optional[int] = 50,
        query: Optional[str] = None,
    ):
        if query is None and not key_names:
            raise ValueError("either `query` or `key_names` must be provided")

        if key_names is not None:
            if not isinstance(key_names, list):
                raise TypeError("key_names must be a list")
            for kn in key_names:
                if not isinstance(kn, str) or not kn:
                    raise ValueError(f"each key_name must be a non-empty string, got {kn!r}")

        if query is not None:
            if not isinstance(query, str) or not query.strip():
                raise ValueError("query must be a non-empty string")

        if result_limit is not None and result_limit < 1:
            raise ValueError("result_limit must be >= 1 or None")

        self._data = data
        self._key_names = key_names or []
        self._result_limit = result_limit
        self._query = query

    def extract(self) -> dict[str, Any]:
        if self._query is not None:
            return {"matches": self._extract_by_query()}
        return {k: self._extract_single(k) for k in self._key_names}

    def _extract_by_query(self) -> List[dict[str, Any]]:
        try:
            expr = ext_parse(self._query)
        except (JsonPathLexerError, JsonPathParserError) as exc:
            raise ValueError(f"invalid jsonpath query {self._query!r}: {exc}") from exc

        matches: List[dict[str, Any]] = [
            {"path": str(m.full_path), "value": m.value}
            for m in expr.find(self._data)
        ]

        if self._result_limit is not None:
            matches = matches[: self._result_limit]

        return matches

    def _extract_single(self, key_name: str) -> List[Any]:
        try:
            expr = _build_jp_expr(key_name)
        except ValueError:
            return []
        vals = [m.value for m in expr.find(self._data)]
        if self._result_limit is not None:
            vals = vals[: self._result_limit]
        return vals


def _build_jp_expr(key_name: str) -> Any:
    if '"' in key_name and "'" in key_name:
        raise ValueError(f"key contains both quote types: {key_name!r}")
    if '"' in key_name:
        expr_str = "$..['" + key_name + "']"
    else:
        expr_str = '$..["' + key_name + '"]'
    try:
        return parse(expr_str)
    except (JsonPathLexerError, JsonPathParserError):
        raise ValueError(f"unable to build jsonpath expression for key: {key_name!r}")


def get_nested_values(data: dict | list, key_name: str) -> List[Any]:
    return KeysExtractor(data, key_names=[key_name], result_limit=None).extract()[key_name]