from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING

from .base import DirectAccessConfigError, TokenHandler, unwrap_secret

if TYPE_CHECKING:
    from relationalai.config.connections.snowflake import SnowflakeConnectionBase


class SnowflakePATHandler(TokenHandler):
    def __init__(self, conn: "SnowflakeConnectionBase"):
        super().__init__(conn)
        self.token_file_path = str(getattr(conn, "token_file_path", "") or "")
        tok_obj = getattr(conn, "token", None)
        self._inline_token = unwrap_secret(tok_obj) or ""

        if not self._inline_token and not self.token_file_path:
            raise DirectAccessConfigError(
                "programmatic_access_token auth requires either `token` or `token_file_path` on the connection."
            )

    @property
    def token(self) -> str:
        if self._inline_token:
            return self._inline_token
        return self._read_token_file()

    def _read_token_file(self) -> str:
        raw = Path(self.token_file_path).read_text(encoding="utf-8").strip()
        # Allow plain token or JSON files produced by different tools.
        if raw.startswith("{"):
            try:
                obj = json.loads(raw)
                if isinstance(obj, dict):
                    for key in ("token", "access_token", "pat", "programmatic_access_token"):
                        if isinstance(obj.get(key), str):
                            return str(obj[key]).strip()
            except Exception:
                pass
        return raw

    def get_session_login_token(self) -> str:
        return self.token

    def get_ingress_token(self, endpoint: str) -> str:
        _ = endpoint
        return self.token


__all__ = ["SnowflakePATHandler"]

