from __future__ import annotations

import logging
import os
import time
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit

from relationalai.client.transports.sql import SqlClient
from relationalai.runtime.execution_summary import SUMMARY_EVENT_SERVICE_ENDPOINT_CACHE_HIT
from relationalai.services._shared.sql_rows import row_first_value, row_get, row_to_mapping
from relationalai.util.persistent_kv import PersistentKeyValueFile

SERVICE_ENDPOINT_CACHE_ENV_VAR = "RAI_ENDPOINT_FILE"
DEFAULT_SERVICE_ENDPOINT_CACHE_PATH = Path(os.path.expanduser("~/.rai/endpoints.json"))
_LOG = logging.getLogger("relationalai.client.execution")


def is_snowflake_notebook() -> bool:
    """Return whether we are running inside a Snowflake Notebook.

    Upstream distinguishes notebook vs local to choose http vs https and the
    boolean parameter passed to `app.service_endpoint(...)`. This repo does not
    have a reliable runtime signal for Snowflake Notebooks, so we conservatively
    treat all environments as non-notebook.
    """
    return False


def default_service_endpoint_cache_path() -> Path:
    return Path(os.environ.get(SERVICE_ENDPOINT_CACHE_ENV_VAR, str(DEFAULT_SERVICE_ENDPOINT_CACHE_PATH)))


def _quote_snowflake_identifier_path(name: str) -> str:
    # Normalize dotted app/object names into Snowflake-safe SQL identifiers while
    # preserving already-quoted or simple uppercase identifiers as-is.
    parts = [part.strip() for part in str(name or "").split(".") if part.strip()]
    quoted: list[str] = []
    for part in parts:
        if part.startswith('"') and part.endswith('"') and len(part) >= 2:
            quoted.append(part)
        elif part.replace("_", "").isalnum() and (part[0].isalpha() or part[0] == "_") and part.upper() == part:
            quoted.append(part)
        else:
            escaped = part.replace('"', '""')
            quoted.append(f'"{escaped}"')
    return ".".join(quoted)


class ServiceEndpointCache(PersistentKeyValueFile):
    """Persistent + in-memory cache for SPCS service endpoints."""

    def __init__(self, path: Path | None = None):
        super().__init__(path or default_service_endpoint_cache_path())

    def get(self, key: str, default: Any = None) -> Any:
        v = super().get(key, default)
        if isinstance(v, dict):
            vv = v.get("value")
            if isinstance(vv, str) and vv.strip():
                return vv.strip()
        if isinstance(v, str) and v.strip():
            return v.strip()
        return default

    def get_cached_at(self, key: str) -> float | None:
        v = super().get(key, None)
        if isinstance(v, dict):
            ts = v.get("cached_at")
            if isinstance(ts, (int, float)):
                return float(ts)
        if v is not None:
            try:
                return float(self.path.stat().st_mtime)
            except Exception:
                return None
        return None

    def set(self, key: str, value: str) -> None:
        v = (value or "").strip()
        if not v:
            return
        super().set(key, {"value": v, "cached_at": time.time()})


def _extract_service_endpoint_host(row: Any) -> str | None:
    """Extract the service endpoint host from a Snowpark `CALL ...` result row.

    Snowpark returns `Row` objects for `.collect()`, but some integration layers
    (tests/mocks) may return plain dicts/tuples. Be liberal in what we accept.
    """
    v = row_get(row, "SERVICE_ENDPOINT", "service_endpoint", "Service_Endpoint", "serviceendpoint")
    if isinstance(v, str) and v.strip():
        return v.strip()
    if not row_to_mapping(row):
        v0 = row_first_value(row)
        if isinstance(v0, str) and v0.strip():
            return v0.strip()
    return None


def _normalize_service_endpoint_base(value: str, *, notebook: bool) -> str:
    raw = str(value or "").strip()
    if not raw:
        raise RuntimeError("Service endpoint result was empty.")

    base = raw if "://" in raw else f"{'http' if notebook else 'https'}://{raw}"
    parts = urlsplit(base)
    if parts.scheme not in {"http", "https"} or not parts.netloc or parts.hostname is None:
        raise RuntimeError(f"Service endpoint returned an invalid URL: {raw!r}.")
    try:
        parts.hostname.encode("idna")
    except UnicodeError as e:
        raise RuntimeError(f"Service endpoint returned an invalid hostname: {parts.hostname!r}.") from e
    return base.rstrip("/")


def resolve_service_endpoint(
    *,
    sql: SqlClient,
    account: str,
    app_name: str,
    cache: ServiceEndpointCache | None = None,
    enforce_update: bool = False,
) -> str:
    """Resolve the Direct Access base URL for a Snowflake SPCS service.

    Matches upstream behavior:
    - Calls `CALL <app>.app.service_endpoint(<bool>)`
    - Uses `https://` for local execution, `http://` for notebook execution.
    - Caches results keyed by (account, app_name).
    """

    acct = (account or "").strip()
    app = (app_name or "").strip()
    if not acct or not app:
        raise ValueError("resolve_service_endpoint requires non-empty account and app_name.")

    cache = cache or ServiceEndpointCache()
    notebook = is_snowflake_notebook()
    # Upstream: `CALL <app>.app.service_endpoint(not is_snowflake_notebook);`
    external = not notebook
    key = f"{acct}.{app}.service_endpoint.external={'true' if external else 'false'}"
    cached = cache.get(key) if cache else None
    cached_at = cache.get_cached_at(key) if cache else None
    if not enforce_update:
        if cached:
            try:
                base = _normalize_service_endpoint_base(str(cached), notebook=notebook)
                _LOG.info(
                    "op=service_endpoint_cache_hit Using cached Direct Access endpoint for app=%s account=%s: %s",
                    app,
                    acct,
                    base,
                    extra={
                        "summary_event": SUMMARY_EVENT_SERVICE_ENDPOINT_CACHE_HIT,
                        "cache_created_at_ts": cached_at,
                    },
                )
                return base
            except RuntimeError as e:
                _LOG.warning(
                    "op=service_endpoint_refresh Discarding invalid cached Direct Access endpoint for app=%s account=%s: %s",
                    app,
                    acct,
                    e,
                )
                if cache:
                    cache.delete(key)

    arg = "TRUE" if external else "FALSE"
    stmt = f"CALL {_quote_snowflake_identifier_path(app)}.app.service_endpoint({arg});"
    rows = sql.collect(stmt, operation="service_endpoint", meta={"app": app, "account": acct})
    if not rows:
        raise RuntimeError(f"Could not retrieve service endpoint for app {app!r}.")

    row0 = rows[0]
    host = _extract_service_endpoint_host(row0)
    if not host:
        raise RuntimeError("Service endpoint result missing SERVICE_ENDPOINT field.")

    base = _normalize_service_endpoint_base(str(host), notebook=notebook)
    if cached:
        try:
            old_base = _normalize_service_endpoint_base(str(cached), notebook=notebook)
        except RuntimeError:
            old_base = None
        if old_base and old_base != base:
            _LOG.info(
                "op=service_endpoint_refresh Updated Direct Access endpoint for app=%s account=%s: %s -> %s",
                app,
                acct,
                old_base,
                base,
            )

    if cache:
        cache.set(key, base)
    return base


__all__ = ["ServiceEndpointCache", "resolve_service_endpoint", "is_snowflake_notebook"]

