"""Tool for JQ-like query and transform on JSON data."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from .base import Tool

_MAX_RESULT_CHARS = 200_000


class JsonQueryTool(Tool):
    """Query and transform JSON data using dot-path expressions."""

    def __init__(self, *, workspace: Path | None = None) -> None:
        self._workspace = workspace

    @property
    def name(self) -> str:
        return "json_query"

    @property
    def description(self) -> str:
        return (
            "Query and transform JSON data. Supports dot-path access, "
            "array indexing, filtering, key selection, and flattening. "
            "Can read JSON from inline text or workspace files."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "data": {
                    "type": "string",
                    "description": "Inline JSON string to query.",
                },
                "file": {
                    "type": "string",
                    "description": "Workspace-relative JSON file to query.",
                },
                "path": {
                    "type": "string",
                    "description": (
                        "Dot-path expression. Examples: "
                        "'items' (key access), 'items[0]' (index), "
                        "'items[*].name' (map), 'items[?price>10]' (filter), "
                        "'keys()' (list keys), 'length()' (count), "
                        "'flatten()' (flatten nested arrays), "
                        "'unique()' (deduplicate), 'sort()' (sort)."
                    ),
                },
                "output_file": {
                    "type": "string",
                    "description": "Save result to this workspace-relative file.",
                },
            },
            "required": ["path"],
        }

    async def execute(self, **kwargs: Any) -> str:
        data_str = kwargs.get("data")
        file_path = kwargs.get("file")
        path_expr = kwargs.get("path", "")
        output_file = kwargs.get("output_file")

        if not data_str and not file_path:
            return json.dumps({"error": "Provide either 'data' (JSON string) or 'file'"})

        if file_path and self._workspace:
            fp = (self._workspace / file_path).resolve()
            if not str(fp).startswith(str(self._workspace.resolve())):
                return json.dumps({"error": "Path outside workspace"})
            if not fp.exists():
                return json.dumps({"error": f"File not found: {file_path}"})
            data_str = fp.read_text(encoding="utf-8")

        try:
            data = json.loads(data_str)
        except (json.JSONDecodeError, TypeError) as exc:
            return json.dumps({"error": f"Invalid JSON: {exc}"})

        try:
            result = self._evaluate(data, path_expr)
        except Exception as exc:
            return json.dumps({"error": f"Query failed: {exc}"})

        if output_file and self._workspace:
            out = (self._workspace / output_file).resolve()
            if not str(out).startswith(str(self._workspace.resolve())):
                return json.dumps({"error": "Output path outside workspace"})
            out.parent.mkdir(parents=True, exist_ok=True)
            out.write_text(json.dumps(result, indent=2), encoding="utf-8")
            return json.dumps({"ok": True, "output_file": str(out), "type": type(result).__name__})

        result_str = json.dumps({"ok": True, "result": result, "type": type(result).__name__})
        if len(result_str) > _MAX_RESULT_CHARS:
            result_str = result_str[:_MAX_RESULT_CHARS] + '..."}'
        return result_str

    def _evaluate(self, data: Any, path: str) -> Any:
        if not path or path == ".":
            return data

        current = data
        tokens = self._tokenize(path)

        for token in tokens:
            current = self._apply_token(current, token)

        return current

    def _tokenize(self, path: str) -> list[str]:
        tokens: list[str] = []
        current = ""
        i = 0
        while i < len(path):
            ch = path[i]
            if ch == ".":
                if current:
                    tokens.append(current)
                    current = ""
            elif ch == "[":
                if current:
                    tokens.append(current)
                    current = ""
                end = path.index("]", i)
                tokens.append(path[i : end + 1])
                i = end
            else:
                current += ch
            i += 1
        if current:
            tokens.append(current)
        return tokens

    def _apply_token(self, data: Any, token: str) -> Any:
        # Functions
        if token == "keys()":
            if isinstance(data, dict):
                return list(data.keys())
            raise ValueError("keys() requires an object")
        if token == "values()":
            if isinstance(data, dict):
                return list(data.values())
            raise ValueError("values() requires an object")
        if token == "length()":
            return len(data)
        if token == "flatten()":
            if isinstance(data, list):
                result = []
                for item in data:
                    if isinstance(item, list):
                        result.extend(item)
                    else:
                        result.append(item)
                return result
            raise ValueError("flatten() requires an array")
        if token == "unique()":
            if isinstance(data, list):
                seen: list = []
                for item in data:
                    if item not in seen:
                        seen.append(item)
                return seen
            raise ValueError("unique() requires an array")
        if token == "sort()":
            if isinstance(data, list):
                return sorted(data, key=lambda x: str(x))
            raise ValueError("sort() requires an array")

        # Array operations
        if token.startswith("[") and token.endswith("]"):
            inner = token[1:-1]

            # Wildcard map
            if inner == "*":
                if not isinstance(data, list):
                    raise ValueError("Wildcard [*] requires an array")
                return data

            # Filter
            if inner.startswith("?"):
                return self._filter(data, inner[1:])

            # Index
            try:
                idx = int(inner)
                return data[idx]
            except (ValueError, IndexError, TypeError):
                raise ValueError(f"Invalid index: {inner}")

        # Key access
        if isinstance(data, dict):
            if token in data:
                return data[token]
            raise ValueError(f"Key not found: {token}")

        # Map over array
        if isinstance(data, list):
            return [item.get(token) if isinstance(item, dict) else None for item in data]

        raise ValueError(f"Cannot access '{token}' on {type(data).__name__}")

    def _filter(self, data: Any, expr: str) -> list:
        if not isinstance(data, list):
            raise ValueError("Filter requires an array")

        import operator
        ops = {">=": operator.ge, "<=": operator.le, ">": operator.gt, "<": operator.lt,
               "==": operator.eq, "!=": operator.ne}

        for op_str, op_fn in sorted(ops.items(), key=lambda x: -len(x[0])):
            if op_str in expr:
                field, value = expr.split(op_str, 1)
                field = field.strip()
                value = value.strip().strip("'\"")
                try:
                    num_val = float(value)
                except ValueError:
                    num_val = None

                results = []
                for item in data:
                    if not isinstance(item, dict) or field not in item:
                        continue
                    item_val = item[field]
                    if num_val is not None and isinstance(item_val, (int, float)):
                        if op_fn(item_val, num_val):
                            results.append(item)
                    elif op_fn(str(item_val), value):
                        results.append(item)
                return results

        raise ValueError(f"Invalid filter expression: {expr}")
