"""
ZoomInfo enrichment client.

Performs OAuth2 token-based calls to ZoomInfo and looks up companies by domain.
Tokens are cached in ~/.g_gremlin/config.json (keys: zoominfo_access_token,
zoominfo_refresh_token, zoominfo_token_expires_at).
"""

from __future__ import annotations

import json
import os
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional

import click

from g_gremlin.auth import resolve_bridge_auth, resolve_runtime_auth
from g_gremlin.base import AuthenticationError, BaseIntegration, IntegrationError, RateLimitError, RetryConfig, TransientError
from g_gremlin.credentials.provider import CredentialProvider, SecretNotFoundError
from g_gremlin.paths import gremlin_path

CONFIG_PATH = gremlin_path("config.json")
TOKEN_URL = "https://api.zoominfo.com/auth/token"
ENRICH_URL = "https://api.zoominfo.com/gtm/data/v1/companies/enrich"
FOUNDRYOPS_DEFAULT_API_BASE = "https://api.foundryops.io"
DEFAULT_RETRY_CONFIG = RetryConfig(max_attempts=4, base_delay=1.0, max_delay=12.0)
JSON_API_HEADERS = {
    "Accept": "application/vnd.api+json",
    "Content-Type": "application/vnd.api+json",
}
DEFAULT_OUTPUT_FIELDS = [
    "name",
    "website",
    "revenue",
    "employeeCount",
    "primaryIndustry",
    "industries",
    "street",
    "city",
    "state",
    "zipCode",
    "country",
]
FIELD_TO_OUTPUT_FIELDS = {
    "name": ("name",),
    "domain": ("website",),
    "revenue": ("revenue",),
    "employee_count": ("employeeCount",),
    "industry": ("primaryIndustry", "industries"),
    "headquarters": ("street", "city", "state", "zipCode", "country"),
}


class ZoomInfoError(RuntimeError):
    """Raised when ZoomInfo operations fail."""


class ZoomInfoAuthError(ZoomInfoError):
    """Raised when access tokens are invalid/expired."""


def _resolve_output_fields(fields: Optional[list[str]]) -> list[str]:
    if not fields:
        return list(DEFAULT_OUTPUT_FIELDS)

    resolved: list[str] = []
    seen: set[str] = set()
    for field in fields:
        for output_field in FIELD_TO_OUTPUT_FIELDS.get(field, ()):
            if output_field in seen:
                continue
            resolved.append(output_field)
            seen.add(output_field)

    if not resolved:
        return list(DEFAULT_OUTPUT_FIELDS)
    return resolved


def _load_config() -> Dict[str, Any]:
    if not CONFIG_PATH.exists():
        return {}
    try:
        return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _save_config(config: Dict[str, Any]) -> None:
    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    CONFIG_PATH.write_text(json.dumps(config, indent=2), encoding="utf-8")


def _parse_expiry(value: Any) -> Optional[float]:
    if value is None:
        return None
    if isinstance(value, (int, float)):
        return float(value)
    raw = str(value).strip()
    if not raw:
        return None
    try:
        return float(raw)
    except ValueError:
        pass
    normalized = raw.replace("Z", "+00:00")
    try:
        parsed = datetime.fromisoformat(normalized)
    except ValueError:
        return None
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.timestamp()


def _resolve_foundryops_api_base(explicit: Optional[str] = None) -> str:
    return (
        explicit
        or os.getenv("FOUNDRYOPS_API_URL")
        or os.getenv("FOUNDRYMATCH_API_URL")
        or FOUNDRYOPS_DEFAULT_API_BASE
    ).rstrip("/")


def _resolve_foundryops_bridge_auth(config: Optional[Dict[str, Any]] = None):
    return resolve_bridge_auth(
        integration="zoominfo",
        config=config,
    )


def _resolve_foundryops_auth_headers(config: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
    return _resolve_foundryops_bridge_auth(config).headers


def resolve_credentials(
    *,
    explicit_credentials: Dict[str, str | None] | None = None,
    credential_provider: CredentialProvider | None = None,
    config: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    resolution = resolve_runtime_auth(
        integration="zoominfo",
        explicit_credentials=explicit_credentials,
        credential_provider=credential_provider,
        config=config,
    )
    creds = dict(resolution.credentials)
    if resolution.mode == "server":
        return {
            "auth_mode": "server",
            "access_token": creds.get("access_token"),
            "token_expires_at": creds.get("token_expires_at"),
        }
    if creds.get("client_id") and creds.get("client_secret"):
        return creds
    raise SecretNotFoundError("Missing credentials for zoominfo")


class ZoomInfoClient(BaseIntegration):
    SERVICE_NAME = "zoominfo"

    def __init__(
        self,
        client_id: Optional[str] = None,
        client_secret: Optional[str] = None,
        credential_provider: CredentialProvider | None = None,
        retry_config: RetryConfig | None = None,
    ) -> None:
        super().__init__(credential_provider=credential_provider, default_retry_config=retry_config or DEFAULT_RETRY_CONFIG)
        self.config = _load_config()
        creds = self._resolve_creds(client_id, client_secret, credential_provider)
        self._auth_mode = str(
            creds.get("auth_mode")
            or self.config.get("zoominfo_auth_mode")
            or "local"
        ).strip().lower()
        self._client_id = str(creds.get("client_id") or "").strip()
        self._client_secret = str(creds.get("client_secret") or "").strip()
        self.config["zoominfo_auth_mode"] = self._auth_mode

    def _resolve_creds(
        self,
        client_id: Optional[str],
        client_secret: Optional[str],
        credential_provider: CredentialProvider | None,
    ) -> Dict[str, Any]:
        explicit_credentials: Dict[str, str | None] | None = None
        if client_id and client_secret:
            explicit_credentials = {
                "client_id": client_id,
                "client_secret": client_secret,
                "auth_mode": "local",
            }

        try:
            creds = resolve_credentials(
                explicit_credentials=explicit_credentials,
                credential_provider=credential_provider,
                config=self.config,
            )
        except SecretNotFoundError as exc:
            raise click.ClickException(
                "ZoomInfo is not configured. "
                "Run 'g-gremlin zoominfo oauth connect' to authenticate."
            ) from exc

        if creds.get("auth_mode") == "server":
            return creds

        if creds.get("access_token"):
            self.config.setdefault("zoominfo_access_token", creds.get("access_token"))
        if creds.get("refresh_token"):
            self.config.setdefault("zoominfo_refresh_token", creds.get("refresh_token"))
        if creds.get("token_expires_at"):
            self.config.setdefault("zoominfo_token_expires_at", creds.get("token_expires_at"))
        if creds.get("client_id") and creds.get("client_secret"):
            return {
                "client_id": creds["client_id"],
                "client_secret": creds["client_secret"],
                "auth_mode": "local",
            }

        raise click.ClickException(
            "ZoomInfo is not configured. "
            "Run 'g-gremlin zoominfo oauth connect' to authenticate."
        )

    # ------------------------------------------------------------------ #
    # Token handling
    # ------------------------------------------------------------------ #
    def _token_valid(self) -> bool:
        expires_at = self.config.get("zoominfo_token_expires_at")
        access_token = self.config.get("zoominfo_access_token")
        if not access_token or not expires_at:
            return False
        expires_ts = _parse_expiry(expires_at)
        if expires_ts is None:
            return False
        return time.time() < expires_ts - 30

    def _store_tokens(self, tokens: Dict[str, Any]) -> None:
        access_token = tokens.get("access_token")
        if not access_token:
            return

        refresh_token = tokens.get("refresh_token") or self.config.get("zoominfo_refresh_token")
        expires_ts = _parse_expiry(tokens.get("expires_at"))
        if expires_ts is None:
            try:
                expires_in = float(tokens.get("expires_in", 3600))
            except Exception:
                expires_in = 3600.0
            expires_ts = time.time() + expires_in

        self.config["zoominfo_access_token"] = access_token
        if self._auth_mode == "server":
            self.config.pop("zoominfo_refresh_token", None)
            self.config.pop("zoominfo_client_id", None)
            self.config.pop("zoominfo_client_secret", None)
        else:
            self.config["zoominfo_refresh_token"] = refresh_token
        self.config["zoominfo_token_expires_at"] = expires_ts
        self.config["zoominfo_auth_mode"] = self._auth_mode
        _save_config(self.config)

    def _server_refresh_token(self) -> str:
        message = "Could not refresh token. Run 'g-gremlin zoominfo oauth connect' to re-authenticate."
        resolution = _resolve_foundryops_bridge_auth(self.config)
        headers = resolution.headers
        if not headers:
            raise ZoomInfoAuthError(message)

        try:
            import requests
        except Exception as exc:
            raise ZoomInfoAuthError(message) from exc

        api_base = resolution.api_base or _resolve_foundryops_api_base(
            str(
                self.config.get("zoominfo_foundryops_api_base")
                or self.config.get("api_base")
                or ""
            )
        )
        params: Dict[str, str] = {}
        workspace_id = str(resolution.workspace_id or "").strip()
        if workspace_id:
            params["workspace_id"] = workspace_id

        try:
            resp = requests.get(
                f"{api_base}/api/v2/integrations/zoominfo/cli/token",
                params=params or None,
                headers=headers,
                timeout=20,
            )
        except Exception as exc:
            raise ZoomInfoAuthError(message) from exc

        if resp.status_code in (401, 403):
            raise ZoomInfoAuthError(message)
        if resp.status_code != 200:
            raise ZoomInfoAuthError(message)

        try:
            payload = resp.json()
        except Exception as exc:
            raise ZoomInfoAuthError(message) from exc

        access_token = str(payload.get("access_token") or "").strip()
        if not access_token:
            raise ZoomInfoAuthError(message)

        self._store_tokens(
            {
                "access_token": access_token,
                "expires_at": payload.get("expires_at"),
                "org_id": payload.get("org_id"),
                "org_name": payload.get("org_name"),
            }
        )
        return access_token

    def _request_tokens(self, refresh_token: Optional[str]) -> Dict[str, Any]:
        if self._auth_mode == "server":
            return {"access_token": self._server_refresh_token()}

        if not self._client_id or not self._client_secret:
            raise ZoomInfoAuthError("ZoomInfo client credentials missing. Reconnect with OAuth.")

        data: Dict[str, Any]
        if refresh_token:
            data = {"grant_type": "refresh_token", "refresh_token": refresh_token}
        else:
            data = {"grant_type": "client_credentials"}

        try:
            resp = self._post(
                TOKEN_URL,
                auth=(self._client_id, self._client_secret),
                json=data,
                allow_statuses=(400,),
                timeout=20,
            )
        except AuthenticationError as exc:
            raise ZoomInfoAuthError("ZoomInfo credentials rejected.") from exc
        except (RateLimitError, IntegrationError, TransientError) as exc:
            raise ZoomInfoError(str(exc)) from exc

        try:
            return resp.json()
        except Exception as exc:  # pragma: no cover - defensive
            raise ZoomInfoError(f"Invalid token response: {exc}") from exc

    def _get_access_token(self) -> str:
        if self._token_valid():
            return self.config["zoominfo_access_token"]

        if self._auth_mode == "server":
            return self._server_refresh_token()

        refresh_token = self.config.get("zoominfo_refresh_token")
        tokens = self._request_tokens(refresh_token if refresh_token else None)
        self._store_tokens(tokens)
        access_token = tokens.get("access_token")
        if not access_token:
            raise ZoomInfoError("ZoomInfo token response missing access_token.")
        return access_token

    # ------------------------------------------------------------------ #
    # Lookup
    # ------------------------------------------------------------------ #
    def _enrich_company(self, token: str, domain: str, fields: Optional[list[str]] = None) -> Dict[str, Any]:
        headers = {"Authorization": f"Bearer {token}", **JSON_API_HEADERS}
        payload = {
            "data": {
                "type": "CompanyEnrich",
                "attributes": {
                    "matchCompanyInput": [{"companyWebsite": domain}],
                    "outputFields": _resolve_output_fields(fields),
                },
            }
        }
        try:
            resp = self._post(
                ENRICH_URL,
                headers=headers,
                json=payload,
            )
        except RateLimitError as exc:
            raise ZoomInfoError(str(exc)) from exc
        except AuthenticationError as exc:
            raise ZoomInfoAuthError(str(exc)) from exc
        except (IntegrationError, TransientError) as exc:
            raise ZoomInfoError(str(exc)) from exc

        try:
            payload = resp.json()
        except Exception as exc:  # pragma: no cover - defensive
            raise ZoomInfoError(f"Invalid ZoomInfo response: {exc}") from exc

        if not isinstance(payload, dict):
            raise ZoomInfoError("Invalid ZoomInfo response: expected object payload.")

        data = payload.get("data")
        if not isinstance(data, list) or not data:
            return {}

        record = data[0]
        if not isinstance(record, dict):
            return {}

        meta = record.get("meta") if isinstance(record.get("meta"), dict) else {}
        match_status = str(meta.get("matchStatus") or "").strip().upper()
        record_type = str(record.get("type") or "").strip()
        if match_status == "LIMIT_EXCEEDED":
            raise ZoomInfoError("ZoomInfo enrich limit exceeded.")
        if record_type == "NoMatch" or match_status in {"NO_MATCH", "INVALID_INPUT"}:
            return {}

        attributes = record.get("attributes")
        if not isinstance(attributes, dict):
            return {}

        company: Dict[str, Any] = dict(attributes)
        if record.get("id") is not None:
            company.setdefault("id", record.get("id"))
        if match_status:
            company.setdefault("matchStatus", match_status)
        return company

    def enrich_company(self, domain: str, fields: Optional[list[str]] = None) -> Dict[str, Any]:
        domain_clean = (domain or "").strip()
        if not domain_clean:
            return {}

        token = self._get_access_token()

        try:
            company = self._enrich_company(token, domain_clean, fields)
        except ZoomInfoAuthError:
            # Refresh tokens and retry once.
            if self._auth_mode == "server":
                token = self._server_refresh_token()
            else:
                tokens = self._request_tokens(self.config.get("zoominfo_refresh_token"))
                self._store_tokens(tokens)
                token = tokens.get("access_token", "")
            company = self._enrich_company(token, domain_clean, fields)

        normalized = _extract_fields(company)
        if fields:
            return {k: normalized.get(k) for k in fields if k in normalized}
        return normalized


def _first_non_empty(values: list[Any]) -> Any:
    for value in values:
        if value is None:
            continue
        if isinstance(value, str):
            if value.strip():
                return value.strip()
            continue
        if isinstance(value, list):
            items = [str(item).strip() for item in value if str(item).strip()]
            if items:
                return items[0]
            continue
        return value
    return None


def _format_headquarters(company: Dict[str, Any]) -> Optional[str]:
    address = {}
    location_match = company.get("locationMatch")
    if isinstance(location_match, dict):
        raw_address = location_match.get("address")
        if isinstance(raw_address, dict):
            address = raw_address

    parts = [
        _first_non_empty([company.get("street"), address.get("street")]),
        _first_non_empty([company.get("city"), address.get("city")]),
        _first_non_empty([company.get("state"), address.get("state")]),
        _first_non_empty([company.get("zipCode"), address.get("zipCode")]),
        _first_non_empty([company.get("country"), address.get("country")]),
    ]
    cleaned = [str(part).strip() for part in parts if str(part or "").strip()]
    if not cleaned:
        return None
    return ", ".join(cleaned)


def _extract_fields(company: Dict[str, Any]) -> Dict[str, Any]:
    if not company:
        return {}
    industry = _first_non_empty(
        [
            company.get("primaryIndustry"),
            company.get("industries"),
            company.get("industry"),
        ]
    )
    headquarters = _first_non_empty(
        [
            _format_headquarters(company),
            company.get("headquarters"),
            company.get("headquartersAddress"),
        ]
    )
    return {
        "name": company.get("name") or company.get("companyName") or company.get("company_name"),
        "domain": company.get("website") or company.get("domain"),
        "revenue": company.get("revenue") or company.get("revenueRange") or company.get("annual_revenue"),
        "employee_count": company.get("employeeCount") or company.get("employee_count"),
        "industry": industry,
        "headquarters": headquarters,
    }


def enrich_company(domain: str, fields: Optional[list[str]] = None) -> Dict[str, Any]:
    """
    Enrich a company by domain using ZoomInfo.
    """
    client = ZoomInfoClient()
    try:
        return client.enrich_company(domain, fields)
    except (ZoomInfoError, ZoomInfoAuthError) as exc:
        raise click.ClickException(str(exc)) from exc


__all__ = [
    "ZoomInfoError",
    "ZoomInfoAuthError",
    "ZoomInfoClient",
    "enrich_company",
]
