from __future__ import annotations

import json
import os
import re
import subprocess
from collections.abc import Set as AbstractSet
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence


Runner = Callable[..., subprocess.CompletedProcess[str]]


class NbCommandError(RuntimeError):
    """Raised when the NocoBase CLI command fails or returns invalid output."""


class NbAuthenticationError(NbCommandError):
    """Raised when NocoBase rejects the configured API credential."""


class NbClient:
    def __init__(
        self,
        *,
        env_name: str = "nb-mcp",
        data_source: str = "main",
        nb_bin: str = "nb",
        api_base_url: str | None = None,
        api_token: str | None = None,
        role: str | None = None,
        cwd: str | Path | None = None,
        timeout: int = 30,
        runner: Runner | None = None,
    ) -> None:
        self.env_name = env_name
        self.data_source = data_source
        expanded_nb_bin = Path(nb_bin).expanduser()
        self.nb_bin = str(expanded_nb_bin) if str(nb_bin).startswith("~") or expanded_nb_bin.is_absolute() else nb_bin
        self.api_base_url = api_base_url.strip().rstrip("/") if api_base_url else None
        self.api_token = api_token
        self.role = role
        self.cwd = Path(cwd).resolve() if cwd else None
        self.timeout = timeout
        self.runner = runner or subprocess.run

    def health(self, resource: str, id_field: str) -> dict[str, Any]:
        self.list_records(resource, fields=[id_field], page=1, page_size=1)
        return {
            "ok": True,
            "env": self.env_name,
            "status_output": "authenticated",
            "stderr": "",
        }

    def list_records(
        self,
        resource: str,
        *,
        filter: Mapping[str, Any] | None = None,
        fields: Sequence[str] | None = None,
        appends: Sequence[str] | None = None,
        except_fields: Sequence[str] | None = None,
        sort: Sequence[str] | None = None,
        page: int | None = None,
        page_size: int | None = None,
        paginate: bool | None = None,
    ) -> dict[str, Any]:
        command = self._resource_command("list", resource)
        self._append_json(command, "--filter", filter)
        self._append_repeated(command, "--fields", fields)
        self._append_repeated(command, "--appends", appends)
        self._append_repeated(command, "--except", except_fields)
        self._append_repeated(command, "--sort", sort)
        self._append_value(command, "--page", page)
        self._append_value(command, "--page-size", page_size)
        if paginate is True:
            command.append("--paginate")
        elif paginate is False:
            command.append("--no-paginate")
        return self._run_json(command)

    def get_record(
        self,
        resource: str,
        filter_by_tk: Any,
        *,
        fields: Sequence[str] | None = None,
        appends: Sequence[str] | None = None,
        except_fields: Sequence[str] | None = None,
    ) -> dict[str, Any]:
        command = self._resource_command("get", resource)
        self._append_json_or_scalar(command, "--filter-by-tk", filter_by_tk)
        self._append_repeated(command, "--fields", fields)
        self._append_repeated(command, "--appends", appends)
        self._append_repeated(command, "--except", except_fields)
        return self._run_json(command)

    def create_record(
        self,
        resource: str,
        values: Mapping[str, Any],
        whitelist: Sequence[str],
    ) -> dict[str, Any]:
        self._validate_scalar_values(values)
        command = self._resource_command("create", resource)
        self._append_json(command, "--values", values)
        self._append_repeated(command, "--whitelist", whitelist)
        return self._run_json(command)

    def update_record(
        self,
        resource: str,
        record_id: Any,
        values: Mapping[str, Any],
        whitelist: Sequence[str],
    ) -> dict[str, Any]:
        self._validate_scalar_values(values)
        command = self._resource_command("update", resource)
        self._append_json_or_scalar(command, "--filter-by-tk", record_id)
        self._append_json(command, "--values", values)
        self._append_repeated(command, "--whitelist", whitelist)
        return self._run_json(command)

    def query_records(
        self,
        resource: str,
        *,
        measures: Sequence[Mapping[str, Any]] | None = None,
        dimensions: Sequence[Mapping[str, Any]] | None = None,
        orders: Sequence[Mapping[str, Any]] | None = None,
        filter: Mapping[str, Any] | None = None,
        having: Mapping[str, Any] | None = None,
        limit: int | None = None,
        offset: int | None = None,
        timezone: str | None = None,
    ) -> dict[str, Any]:
        command = self._resource_command("query", resource)
        self._append_json(command, "--measures", measures)
        self._append_json(command, "--dimensions", dimensions)
        self._append_json(command, "--orders", orders)
        self._append_json(command, "--filter", filter)
        self._append_json(command, "--having", having)
        self._append_value(command, "--limit", limit)
        self._append_value(command, "--offset", offset)
        self._append_value(command, "--timezone", timezone)
        return self._run_json(command)

    def get_collection(
        self,
        collection_name: str,
        *,
        appends: Sequence[str] | None = None,
        except_fields: Sequence[str] | None = None,
    ) -> dict[str, Any]:
        command = self._data_modeling_command("collections", "get")
        command.extend(["--filter-by-tk", collection_name])
        self._append_repeated(command, "--appends", appends)
        self._append_repeated(command, "--except", except_fields)
        return self._run_json(command)

    def list_collection_fields(
        self,
        collection_name: str,
        *,
        fields: Sequence[str] | None = None,
        appends: Sequence[str] | None = None,
        except_fields: Sequence[str] | None = None,
        page: int | None = None,
        page_size: int | None = None,
        paginate: bool | None = None,
    ) -> dict[str, Any]:
        command = self._data_modeling_command("collections", "fields", "list")
        command.extend(["--collection-name", collection_name])
        self._append_repeated(command, "--fields", fields)
        self._append_repeated(command, "--appends", appends)
        self._append_repeated(command, "--except", except_fields)
        self._append_value(command, "--page", page)
        self._append_value(command, "--page-size", page_size)
        if paginate is True:
            command.append("--paginate")
        elif paginate is False:
            command.append("--no-paginate")
        return self._run_json(command)

    def _resource_command(self, action: str, resource: str) -> list[str]:
        command = [
            self.nb_bin,
            "api",
            "resource",
            action,
            "--resource",
            resource,
            "--env",
            self.env_name,
            "--yes",
            "--json-output",
        ]
        if self.data_source:
            command.extend(["--data-source", self.data_source])
        if self.api_base_url:
            command.extend(["--api-base-url", self.api_base_url])
        if self.role:
            command.extend(["--role", self.role])
        if self.api_token:
            command.extend(["--token", self.api_token])
        return command

    def _data_modeling_command(self, *parts: str) -> list[str]:
        command = [
            self.nb_bin,
            "api",
            "data-modeling",
            *parts,
            "--env",
            self.env_name,
            "--yes",
            "--json-output",
        ]
        if self.api_base_url:
            command.extend(["--api-base-url", self.api_base_url])
        if self.role:
            command.extend(["--role", self.role])
        if self.api_token:
            command.extend(["--token", self.api_token])
        return command

    def _run_json(self, command: list[str]) -> dict[str, Any]:
        result = self._run_raw(command[1:], full_command=command)
        if result.returncode != 0:
            error_type = NbAuthenticationError if self._is_authentication_failure(result) else NbCommandError
            raise error_type(self._format_error(command, result))
        try:
            parsed = self._parse_json_object(result.stdout)
        except json.JSONDecodeError as error:
            raise NbCommandError(
                f"nb returned invalid JSON: {error.msg}\nstdout={self._redact(result.stdout.strip())}"
            ) from error
        return parsed

    @staticmethod
    def _parse_json_object(output: str) -> dict[str, Any]:
        stripped = output.strip()
        try:
            parsed = json.loads(stripped)
        except json.JSONDecodeError:
            parsed = None
        if parsed is not None:
            if not isinstance(parsed, dict):
                raise NbCommandError("nb returned JSON that is not an object")
            return parsed

        decoder = json.JSONDecoder()
        candidates: list[tuple[int, int, dict[str, Any]]] = []
        for start, character in enumerate(output):
            if character != "{":
                continue
            try:
                value, end = decoder.raw_decode(output, start)
            except json.JSONDecodeError:
                continue
            if isinstance(value, dict):
                candidates.append((start, end, value))

        outer_candidates = [
            candidate
            for candidate in candidates
            if not any(
                other_start <= candidate[0]
                and other_end >= candidate[1]
                and (other_start, other_end) != (candidate[0], candidate[1])
                for other_start, other_end, _ in candidates
            )
        ]
        if len(outer_candidates) != 1:
            raise json.JSONDecodeError("could not locate one complete JSON object", output, 0)
        return outer_candidates[0][2]

    def _run_raw(
        self,
        args: list[str],
        *,
        full_command: list[str] | None = None,
    ) -> subprocess.CompletedProcess[str]:
        command = full_command or [self.nb_bin, *args]
        env = os.environ.copy()
        env["NB_CLI_LOG_DISABLED"] = "1"
        env["NB_SKIP_STARTUP_UPDATE"] = "1"
        self._prepend_nb_bin_dir_to_path(env)
        return self.runner(
            command,
            cwd=str(self.cwd) if self.cwd else None,
            env=env,
            timeout=self.timeout,
            capture_output=True,
            text=True,
            check=False,
        )

    def _prepend_nb_bin_dir_to_path(self, env: dict[str, str]) -> None:
        nb_path = Path(self.nb_bin).expanduser()
        if not nb_path.is_absolute():
            return
        nb_dir = str(nb_path.parent)
        existing_path = env.get("PATH", "")
        parts = [part for part in existing_path.split(os.pathsep) if part]
        if nb_dir in parts:
            return
        env["PATH"] = nb_dir + (os.pathsep + existing_path if existing_path else "")

    def _format_error(self, command: list[str], result: subprocess.CompletedProcess[str]) -> str:
        return (
            f"nb command failed with exit code {result.returncode}: {self._redact_command(command)}\n"
            f"stdout={self._redact(result.stdout.strip())}\n"
            f"stderr={self._redact(result.stderr.strip())}"
        )

    @staticmethod
    def _is_authentication_failure(result: subprocess.CompletedProcess[str]) -> bool:
        output = f"{result.stdout}\n{result.stderr}"
        if re.search(r"\bstatus\s+401\b", output, flags=re.IGNORECASE):
            return True
        upper_output = output.upper()
        return any(
            code in upper_output
            for code in (
                "EMPTY_TOKEN",
                "INVALID_TOKEN",
                "EXPIRED_TOKEN",
                "BLOCKED_TOKEN",
                "EXPIRED_SESSION",
                "NOT_EXIST_USER",
            )
        )

    def _redact_command(self, command: Sequence[str]) -> str:
        redacted: list[str] = []
        skip_next = False
        for token in command:
            if skip_next:
                redacted.append("<redacted>")
                skip_next = False
                continue
            redacted.append(token)
            if token == "--token":
                skip_next = True
        return " ".join(redacted)

    def _redact(self, value: str) -> str:
        redacted = value
        if self.api_token:
            redacted = redacted.replace(self.api_token, "<redacted>")
        return redacted

    @staticmethod
    def _append_json(command: list[str], flag: str, value: Any) -> None:
        if value is None or value == {}:
            return
        command.extend([flag, json.dumps(value, ensure_ascii=False, separators=(",", ":"))])

    @staticmethod
    def _append_json_or_scalar(command: list[str], flag: str, value: Any) -> None:
        if value is None:
            return
        if isinstance(value, (list, dict)):
            command.extend([flag, json.dumps(value, ensure_ascii=False, separators=(",", ":"))])
        else:
            command.extend([flag, str(value)])

    @staticmethod
    def _append_repeated(command: list[str], flag: str, values: Sequence[str] | None) -> None:
        for value in values or []:
            command.extend([flag, str(value)])

    @staticmethod
    def _append_value(command: list[str], flag: str, value: Any) -> None:
        if value is None or value == "":
            return
        command.extend([flag, str(value)])

    @staticmethod
    def _validate_scalar_values(values: Mapping[str, Any]) -> None:
        if not values:
            raise ValueError("values must not be empty")
        for field, value in values.items():
            if isinstance(value, (Mapping, list, AbstractSet, tuple)):
                raise ValueError(f"values[{field!r}] must be scalar")