from __future__ import annotations

from typing import Any, Dict, List, Tuple, Optional
import os
import time
import httpx


class _TokenBucket:
    """Simple token bucket for rate limiting (best-effort, thread-safe-enough for our use).
    capacity tokens, refill per second.
    """

    def __init__(self, capacity: int, refill_rate_per_sec: float) -> None:
        self.capacity = max(1, capacity)
        self.tokens = float(self.capacity)
        self.refill = float(refill_rate_per_sec)
        self._last = time.time()

    def wait_for_token(self) -> None:
        now = time.time()
        elapsed = max(0.0, now - self._last)
        self._last = now
        self.tokens = min(self.capacity, self.tokens + elapsed * self.refill)
        if self.tokens < 1.0:
            # sleep until 1 token is available
            need = 1.0 - self.tokens
            delay = need / max(0.001, self.refill)
            time.sleep(delay)
            self.tokens = 0.0
            self._last = time.time()
        else:
            self.tokens -= 1.0


class ApolloAdapter:
    """Apollo enrichment adapter (stubbed calls, real retries/backoff scaffolding).

    Configure with env:
      APOLLO_API_KEY
      APOLLO_BASE_URL (optional; default https://api.apollo.io/v1)
      APOLLO_RPS (optional; default 3)
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        rps: Optional[int] = None,
    ) -> None:
        self.api_key = api_key or os.getenv("APOLLO_API_KEY", "")
        self.base_url = (
            base_url or os.getenv("APOLLO_BASE_URL", "https://api.apollo.io/v1")
        ).rstrip("/")
        self.rps = int(os.getenv("APOLLO_RPS", str(rps or 3)))
        self._bucket = _TokenBucket(
            capacity=max(1, self.rps * 2), refill_rate_per_sec=float(self.rps)
        )
        self._client = httpx.Client(timeout=20.0)

    def _headers(self) -> Dict[str, str]:
        return {"Authorization": f"Bearer {self.api_key}", "Accept": "application/json"}

    def available(self) -> bool:
        return bool(self.api_key)

    def enrich_batch(
        self,
        headers: List[str],
        rows: List[List[Any]],
        fields: List[str],
    ) -> Tuple[List[List[Any]], float]:
        """Call Apollo in batches. Returns (enriched_rows, credits_used).
        Expected Apollo response shape:
          { "headers": [...], "rows": [[...], ...], "credits_used": 12.3 }
        """
        if not self.available():
            return rows, 0.0

        payload: Dict[str, Any] = {"headers": headers, "rows": rows, "fields": fields}

        self._bucket.wait_for_token()
        try:
            resp = self._client.post(
                f"{self.base_url}/enrich/batch",
                json=payload,
                headers=self._headers(),
            )
            resp.raise_for_status()
            data = resp.json() or {}
        except httpx.HTTPError:
            # Best-effort fallback: pass-through
            return rows, 0.0

        ap_headers = data.get("headers") or headers
        ap_rows = data.get("rows") or rows
        credits = float(data.get("credits_used") or 0.0)

        mapped_rows, _ = _map_apollo_rows(headers, ap_headers, ap_rows)
        return mapped_rows or rows, credits


def _map_apollo_rows(
    orig_headers: List[str], ap_headers: List[str], ap_rows: List[List[Any]]
) -> Tuple[List[List[Any]], List[str]]:
    """Align Apollo rows to original column order and append any new fields at the end."""
    ap_idx: Dict[str, int] = {str(h or "").strip(): i for i, h in enumerate(ap_headers)}
    out_headers: List[str] = list(orig_headers)
    # Append new fields once
    for h in ap_headers:
        if h not in out_headers:
            out_headers.append(h)

    out_rows: List[List[Any]] = []
    for r in ap_rows:
        base = [""] * len(out_headers)
        # fill original columns
        for i, h in enumerate(orig_headers):
            j = ap_idx.get(h)
            base[i] = r[j] if (j is not None and j < len(r)) else ""
        # fill appended columns
        for h in ap_headers:
            if h not in orig_headers:
                j = ap_idx.get(h)
                base[out_headers.index(h)] = (
                    r[j] if (j is not None and j < len(r)) else ""
                )
        out_rows.append(base)
    return out_rows, out_headers
