# liberal_alpha/client.py
from __future__ import annotations

import base64
import os
import time
import math
import json
import logging
import datetime as dt
import random
import sys
import shutil
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, Tuple, Callable

import requests

logger = logging.getLogger(__name__)

try:
    import pandas as pd
except Exception:  # pragma: no cover
    pd = None  # type: ignore

try:
    from zoneinfo import ZoneInfo
except Exception:  # pragma: no cover
    ZoneInfo = None  # type: ignore

from .crypto import decrypt_data_entry_payload
from .proto_utils import parse_length_prefixed_messages
from .utils import get_user_wallet_address

# ---- endpoints ----
DEFAULT_API_BASE = os.getenv("LIBALPHA_API_BASE", "https://api.librealpha.com").rstrip("/")
# protobuf file download (length-prefixed)
ENTRY_DOWNLOAD_LINKS_PATH = "/api/entries/download-links"  # GET
SUBSCRIPTIONS_PATH = "/api/subscriptions"  # GET (JWT or API key)

# upload endpoints (api-key auth)
# ✅ V2 additions (do not change existing logic; only add new capabilities)
RECORDS_PATH = "/api/records"
SUBSCRIPTIONS_CANDIDATE_PATHS = [
    "/api/subscriptions",
    "/api/users/subscriptions",
    "/api/users/me/subscriptions",
    "/api/records/subscriptions",
]
WS_V1_PATH = os.getenv("LIBALPHA_WS_V1_PATH", "/ws/data")
WS_V2_PATH = os.getenv("LIBALPHA_WS_V2_PATH", "/ws/data/v2")


# ----------------------------
# Exceptions (keep simple)
# ----------------------------
class ConfigurationError(Exception):
    pass


class RequestError(Exception):
    pass


# ----------------------------
# helpers
# ----------------------------
def _ensure_pandas():
    if pd is None:
        raise ImportError("pandas is required. Please `pip install pandas`.")


def _normalize_encryption_key(key: Optional[str]) -> Optional[str]:
    if key is None:
        return None
    key = str(key).strip()
    return key or None


def _utc_iso_z(ts: Optional[dt.datetime] = None) -> str:
    if ts is None:
        ts = dt.datetime.now(dt.timezone.utc)
    ts = ts.astimezone(dt.timezone.utc)
    return ts.replace(microsecond=(ts.microsecond // 1000) * 1000).isoformat().replace("+00:00", "Z")


def _parse_iso_any(s: str) -> Optional[dt.datetime]:
    if not s or not isinstance(s, str):
        return None
    s = s.strip()
    try:
        if s.endswith("Z"):
            s = s[:-1] + "+00:00"
        return dt.datetime.fromisoformat(s)
    except Exception:
        return None


def _normalize_tz(tz_info: Union[dt.tzinfo, str, int, float, None]) -> dt.tzinfo:
    """
    Accept:
      - tzinfo instance
      - "Asia/Singapore"
      - "+8" / "-4"
      - 8 / -4 / 5.5 etc
    """
    if tz_info is None:
        return dt.timezone.utc

    if isinstance(tz_info, dt.tzinfo):
        return tz_info

    if isinstance(tz_info, (int, float)):
        hours = float(tz_info)
        minutes = int(round((hours - math.trunc(hours)) * 60))
        return dt.timezone(dt.timedelta(hours=int(math.trunc(hours)), minutes=minutes))

    if isinstance(tz_info, str):
        s = tz_info.strip()
        # "+8" "-4" "5.5"
        if (s.startswith(("+", "-")) and s[1:].replace(".", "", 1).isdigit()) or s.replace(".", "", 1).isdigit():
            return _normalize_tz(float(s))
        # "Asia/xxx"
        if ZoneInfo is None:
            raise ImportError("zoneinfo not available. Use tz offset int like 8 / -4.")
        return ZoneInfo(s)

    raise TypeError(f"Unsupported tz_info type: {type(tz_info)}")


def _guess_timestamp_ms(value: Any) -> Optional[int]:
    if value is None:
        return None
    try:
        iv = int(value)
    except Exception:
        return None
    # seconds vs ms heuristic
    if iv < 10_000_000_000:
        return iv * 1000
    return iv


def _guess_timestamp_us(value: Any) -> Optional[int]:
    """
    Heuristic for unix timestamps:
      - seconds  (1e9)  -> * 1e6
      - millis   (1e12) -> * 1e3
      - micros   (1e15) -> as-is
    """
    if value is None:
        return None
    try:
        iv = int(value)
    except Exception:
        return None
    # seconds
    if iv < 100_000_000_000:
        return iv * 1_000_000
    # milliseconds
    if iv < 100_000_000_000_000:
        return iv * 1_000
    # microseconds
    return iv


def _dt_to_us(ts: dt.datetime) -> int:
    if not isinstance(ts, dt.datetime):
        raise TypeError("ts must be datetime")
    if ts.tzinfo is None:
        # treat naive as UTC for backward-compat
        ts = ts.replace(tzinfo=dt.timezone.utc)
    ts = ts.astimezone(dt.timezone.utc)
    return int(ts.timestamp() * 1_000_000)


def _normalize_time_range_us(
    start: Union[int, float, str, dt.datetime],
    end: Union[int, float, str, dt.datetime],
) -> Tuple[int, int]:
    """
    Normalize inputs to (start_us, end_us) inclusive.
    Accept:
      - datetime
      - int/float: seconds/ms/us heuristic
      - str: ISO8601
    """
    def to_us(v: Union[int, float, str, dt.datetime]) -> Optional[int]:
        if isinstance(v, dt.datetime):
            return _dt_to_us(v)
        if isinstance(v, str):
            dd = _parse_iso_any(v)
            return _dt_to_us(dd) if dd else None
        if isinstance(v, (int, float)):
            return _guess_timestamp_us(int(v))
        return None

    s = to_us(start)
    e = to_us(end)
    if s is None or e is None:
        raise ValueError("start/end must be datetime, ISO string, or unix timestamp (sec/ms/us)")
    if e < s:
        raise ValueError("end must be >= start")
    return int(s), int(e)


def _split_range_us_24h(start_us: int, end_us: int) -> List[Tuple[int, int]]:
    """
    Split [start_us, end_us] into chunks where each chunk <= 24h.
    """
    max_span = int(dt.timedelta(hours=24).total_seconds() * 1_000_000)
    out: List[Tuple[int, int]] = []
    cur = int(start_us)
    while cur <= int(end_us):
        chunk_end = min(int(end_us), cur + max_span - 1)
        out.append((cur, chunk_end))
        cur = chunk_end + 1
    return out


def _normalize_download_links(payload: Any) -> List[Dict[str, Any]]:
    """
    backend: { status, data: { links: [ {path, download_url, minute}, ... ] } }
    """
    if not isinstance(payload, dict):
        return []
    data = payload.get("data")
    if not isinstance(data, dict):
        return []
    links = data.get("links")
    if not isinstance(links, list):
        return []
    return [x for x in links if isinstance(x, dict)]


class _SingleLineProgress:
    """
    Single-line progress bar rendered with carriage returns (no multi-line spam).
    """

    def __init__(self, *, enabled: bool):
        self.enabled = bool(enabled)
        self._last_len = 0
        # stdout tends to behave better than stderr in IDE/Jupyter output panels.
        self._stream = sys.stdout
        # Prefer ANSI clear-line when supported; still keep padding fallback.
        self._use_ansi = os.getenv("LIBALPHA_PROGRESS_ANSI", "1") != "0"
        # Notebook (ipykernel) often does not honor \r overwrite; use display_id update if available.
        self._ipython_handle = None
        if self.enabled:
            try:
                if "ipykernel" in sys.modules:
                    from IPython.display import display  # type: ignore

                    self._ipython_handle = display("", display_id=True)
            except Exception:
                self._ipython_handle = None
        # Throttle renders to reduce output pressure in busy loops.
        self._min_interval_s = float(os.getenv("LIBALPHA_PROGRESS_MIN_INTERVAL_S", "0.12"))
        self._last_render_ts = 0.0
        self._last_text = ""

    def update(self, text: str) -> None:
        if not self.enabled:
            return
        s = str(text).replace("\r", " ").replace("\n", " ")
        now = time.time()
        if s == self._last_text:
            return
        if self._min_interval_s > 0 and (now - self._last_render_ts) < self._min_interval_s:
            # Best-effort throttle; keep the latest text for next tick.
            self._last_text = s
            return

        # Keep line short to avoid terminal auto-wrap (wrap makes it look like multi-line spam).
        try:
            term_w = int(shutil.get_terminal_size(fallback=(120, 20)).columns)
        except Exception:
            term_w = 120
        term_w = max(40, min(term_w, 240))
        max_len = term_w - 2
        if len(s) > max_len:
            if max_len >= 10:
                s = s[: max_len - 3] + "..."
            else:
                s = s[:max_len]

        pad = ""
        if self._last_len > len(s):
            pad = " " * (self._last_len - len(s))
        # Use carriage return + clear-line + padding fallback.
        # Some environments ignore \r; ANSI clear-line improves odds of single-line rendering.
        if self._ipython_handle is not None:
            try:
                self._ipython_handle.update(s)
            except Exception:
                self._ipython_handle = None
        if self._ipython_handle is None:
            prefix = "\r"
            if self._use_ansi:
                prefix = "\r\x1b[2K"
            self._stream.write(prefix + s + pad)
            self._stream.flush()
        self._last_len = len(s)
        self._last_render_ts = now
        self._last_text = s

    def finish(self) -> None:
        if not self.enabled:
            return
        if self._ipython_handle is None:
            self._stream.write("\n")
            self._stream.flush()
        self._last_len = 0
        self._last_text = ""


def _format_progress_bar(*, pct: float, width: int = 28) -> str:
    p = max(0.0, min(100.0, float(pct)))
    filled = int(round((p / 100.0) * width))
    filled = max(0, min(width, filled))
    return "[" + ("#" * filled) + ("-" * (width - filled)) + f"] {p:6.2f}%"


def _format_bytes(n: int) -> str:
    n = int(n)
    if n < 0:
        n = 0
    units = ["B", "KB", "MB", "GB", "TB"]
    x = float(n)
    u = 0
    while x >= 1024.0 and u < len(units) - 1:
        x /= 1024.0
        u += 1
    if u == 0:
        return f"{int(x)}{units[u]}"
    return f"{x:.1f}{units[u]}"


def _infer_entry_type_from_path(path: Any) -> Optional[str]:
    if not path:
        return None
    p = str(path).lower()
    if "alpha_record_" in p:
        return "alpha"
    if "record_" in p:
        return "data"
    return None


def _local_yyyymmdd(ts_ms: int, tz: dt.tzinfo) -> int:
    d = dt.datetime.fromtimestamp(ts_ms / 1000.0, tz=dt.timezone.utc).astimezone(tz).date()
    return d.year * 10000 + d.month * 100 + d.day


def _cursor_end_of_day_utc(yyyymmdd: int, tz: dt.tzinfo) -> str:
    """
    Given local date (yyyymmdd) in tz, return UTC ISO cursor string at that day's 23:59:59.999 local time.
    """
    y = int(yyyymmdd) // 10000
    m = (int(yyyymmdd) // 100) % 100
    d = int(yyyymmdd) % 100

    local_dt = dt.datetime(y, m, d, 23, 59, 59, 999000, tzinfo=tz)
    utc_dt = local_dt.astimezone(dt.timezone.utc)
    return utc_dt.isoformat().replace("+00:00", "Z")


def _ensure_0x(hex_str: str) -> str:
    if not isinstance(hex_str, str):
        hex_str = str(hex_str)
    hex_str = hex_str.strip()
    if not hex_str.startswith("0x"):
        return "0x" + hex_str
    return hex_str


def _strip_0x(hex_str: str) -> str:
    if not isinstance(hex_str, str):
        hex_str = str(hex_str)
    hex_str = hex_str.strip()
    if hex_str.startswith("0x"):
        return hex_str[2:]
    return hex_str


def _is_nan(v: Any) -> bool:
    return isinstance(v, float) and math.isnan(v)


# ✅ V2 additions: ws url helper (no impact on existing logic)
def _to_ws_url(api_base: str) -> str:
    """
    https://api.xxx -> wss://api.xxx
    http://localhost:8080 -> ws://localhost:8080
    """
    b = (api_base or "").strip()
    if b.startswith("https://"):
        return "wss://" + b[len("https://") :]
    if b.startswith("http://"):
        return "ws://" + b[len("http://") :]
    return b


# ----------------------------
# Client
# ----------------------------
class LiberalAlphaClient:
    """
    API client (no local runner required).

    Auth model:
      - API key (X-API-Key) for all HTTP APIs used by this SDK.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        api_base: str = DEFAULT_API_BASE,
        timeout: int = 30,
        # ✅ V2 additions: runner compat (optional; defaults keep old behavior)
        host: Optional[str] = None,
        port: Optional[int] = None,
        base_url: Optional[str] = None,
        # HTTP behavior: whether to respect environment / system proxy settings (WinINET)
        # In some corporate networks, system proxy can break TLS for Python requests (EOF).
        trust_env: Optional[bool] = None,
    ):
        # ✅ keep existing behavior: api_base default stays, but allow base_url override for v2 script
        if base_url:
            api_base = base_url

        self.api_key = api_key or os.getenv("LIBALPHA_API_KEY")
        self.api_base = (api_base or DEFAULT_API_BASE).rstrip("/")
        self.timeout = timeout

        # ----------------------------
        # HTTP session (for connection reuse + optional proxy bypass)
        # ----------------------------
        if trust_env is None:
            # Default: respect env/system proxy, but allow disabling via env var.
            # LIBALPHA_HTTP_TRUST_ENV=0 => disable
            trust_env = os.getenv("LIBALPHA_HTTP_TRUST_ENV", "1") != "0"
        self._http_trust_env = bool(trust_env)
        self._http = requests.Session()
        self._http.trust_env = self._http_trust_env

        # ✅ runner fields (not used unless you call send_data)
        self.runner_host = host or os.getenv("LIBALPHA_RUNNER_HOST", "127.0.0.1")
        self.runner_port = int(port or os.getenv("LIBALPHA_RUNNER_PORT", "8128"))

        # ----------------------------
        # Minimal HTTP retry policy
        # ----------------------------
        # Retry ONCE for transient backend/network issues:
        # - status: 502/503/504
        # - network: timeout / connection error
        self._http_retry_503 = int(os.getenv("LIBALPHA_HTTP_RETRY_503", "1"))
        self._http_retry_timeout = int(os.getenv("LIBALPHA_HTTP_RETRY_TIMEOUT", "1"))
        self._http_retry_sleep = float(os.getenv("LIBALPHA_HTTP_RETRY_SLEEP", "1.0"))
        self._http_retry_jitter = float(os.getenv("LIBALPHA_HTTP_RETRY_JITTER", "0.5"))

    def _is_retryable_status(self, code: int) -> bool:
        return code in (502, 503, 504)

    def _sleep_before_retry(self, attempt: int) -> None:
        # attempt: 0 for first retry
        time.sleep(self._http_retry_sleep + random.random() * self._http_retry_jitter)

    # ----------------------------
    # HTTP helpers
    # ----------------------------
    def _headers_api_key(self) -> Dict[str, str]:
        h = {"Accept": "application/json"}
        if self.api_key:
            h["X-API-Key"] = self.api_key
        return h

    def _request(
        self,
        method: str,
        path: str,
        *,
        params: Optional[Dict[str, Any]] = None,
        json_body: Optional[Dict[str, Any]] = None,
        auth: str = "api_key",  # "api_key" or "none"
        _retry: bool = True,
    ) -> Any:
        url = f"{self.api_base}{path}"

        if auth == "api_key":
            headers = self._headers_api_key()
            if not self.api_key:
                raise ConfigurationError("Missing api_key (LIBALPHA_API_KEY).")
        else:
            headers = {"Accept": "application/json"}

        max_status_retries = self._http_retry_503 if _retry else 0
        max_net_retries = self._http_retry_timeout if _retry else 0

        status_retry_used = 0
        net_retry_used = 0

        while True:
            try:
                resp = self._http.request(
                    method=method,
                    url=url,
                    headers=headers,
                    params=params,
                    json=json_body,
                    timeout=self.timeout,
                )
            except (requests.Timeout, requests.ConnectionError) as e:
                if net_retry_used < max_net_retries:
                    net_retry_used += 1
                    logger.warning(
                        "HTTP network/timeout error (%s %s): %s. retry=%s/%s",
                        method,
                        url,
                        e,
                        net_retry_used,
                        max_net_retries,
                    )
                    self._sleep_before_retry(net_retry_used - 1)
                    continue
                raise RequestError(f"HTTP request failed: {e}") from e
            except requests.RequestException as e:
                raise RequestError(f"HTTP request failed: {e}") from e

            # retry once for 502/503/504
            if self._is_retryable_status(resp.status_code) and status_retry_used < max_status_retries:
                status_retry_used += 1
                logger.warning(
                    "HTTP %s from %s %s: %s. retry=%s/%s",
                    resp.status_code,
                    method,
                    url,
                    resp.text.strip()[:200],
                    status_retry_used,
                    max_status_retries,
                )
                self._sleep_before_retry(status_retry_used - 1)
                continue

            if resp.status_code >= 400:
                raise RequestError(f"HTTP {resp.status_code}: {resp.text.strip()}")

            ct = resp.headers.get("Content-Type", "")
            if "application/json" in ct:
                try:
                    return resp.json()
                except Exception:
                    return resp.text

            try:
                return resp.json()
            except Exception:
                return resp.text

    # ----------------------------
    # Records endpoint (optional)
    # ----------------------------
    def my_records(self) -> Any:
        return self._request("GET", "/api/records", auth="api_key")

    # ✅ V2 additions: subscriptions (does not alter any existing behavior)
    def my_subscriptions(self) -> Any:
        last_err: Optional[Exception] = None
        for path in SUBSCRIPTIONS_CANDIDATE_PATHS:
            try:
                return self._request("GET", path, auth="api_key")
            except Exception as e:
                last_err = e
                continue
        raise RequestError(f"Cannot fetch subscriptions. Tried paths={SUBSCRIPTIONS_CANDIDATE_PATHS}. Last error={last_err}")

    # ----------------------------
    # Public API: download_history_data() (protobuf length-prefixed)
    # ----------------------------
    def download_history_data(
        self,
        *,
        record_id: int,
        symbols: List[str],
        start: Union[int, float, str, dt.datetime],
        end: Union[int, float, str, dt.datetime],
        encryption_key: Optional[str] = None,
    ):
        """
        Download historical data by time range using /api/entries/download-links.

        Backend constraint: end-start must be <= 24h, so this method auto-splits
        the requested range and merges results locally.

        When data is encrypted (features empty, encrypted_payload present), pass
        encryption_key to decrypt; plaintext is uint32 count + float64[count].

        Params:
          - record_id: int
          - symbols: list[str] (symbol names; [] => auto fetch all target_symbols)
          - start/end: datetime | ISO string | unix timestamp (sec/ms/us)
          - encryption_key: optional, used to decrypt encrypted_payload when features are empty
        Returns:
            pandas.DataFrame
        """
        _ensure_pandas()

        if record_id <= 0:
            raise ValueError("record_id must be positive int")
        if not self.api_key:
            raise ConfigurationError("Missing api_key (LIBALPHA_API_KEY). download_history_data needs auth.")

        syms_in = [str(s).strip() for s in (symbols or []) if str(s).strip()]
        if not syms_in:
            # Auto fetch all target_symbols (symbol names) from /api/records/user-records
            ur = self._request("GET", "/api/records/user-records", auth="api_key")
            recs = (ur.get("data") or {}).get("records") if isinstance(ur, dict) else None
            if not isinstance(recs, list):
                recs = []
            rec = next((x for x in recs if isinstance(x, dict) and int(x.get("id") or 0) == int(record_id)), None)
            ts = (rec or {}).get("target_symbols") if isinstance(rec, dict) else None
            if not isinstance(ts, list):
                ts = []
            for item in ts:
                if isinstance(item, dict) and isinstance(item.get("symbol"), str) and item.get("symbol").strip():
                    syms_in.append(item["symbol"].strip())
            if not syms_in:
                raise ConfigurationError(
                    "symbols is empty and record has no target_symbols; please pass symbols explicitly."
                )

        start_us, end_us = _normalize_time_range_us(start, end)
        chunks = _split_range_us_24h(start_us, end_us)

        try:
            from .proto import data_entry_pb2
        except Exception as e:
            raise ConfigurationError(
                "Cannot import data_entry_pb2. Ensure SDK package includes liberal_alpha/proto/data_entry_pb2.py "
                "and protobuf>=5.29.0 is installed."
            ) from e

        # data / alpha 统一使用 DataEntry（features 为 repeated double）
        msg_cls = data_entry_pb2.DataEntry

        # Use the same session to keep consistent trust_env behavior (proxy/no-proxy)
        session = self._http
        dl_timeout = max(int(self.timeout), 60)

        rows: List[Dict[str, Any]] = []
        progress = _SingleLineProgress(enabled=os.getenv("LIBALPHA_DOWNLOAD_PROGRESS", "1") != "0")
        if progress.enabled:
            progress.update(f"Downloading {_format_progress_bar(pct=0.0)} preparing...")

        # Collect all links first so we can show an overall progress bar.
        all_items: List[Dict[str, Any]] = []
        total_link_calls = max(len(syms_in), 1) * max(len(chunks), 1)
        link_call_idx = 0
        for sym in syms_in:
            for s_us, e_us in chunks:
                link_call_idx += 1
                if progress.enabled:
                    pct_stage = (link_call_idx / max(total_link_calls, 1)) * 100.0
                    progress.update(
                        f"Downloading {_format_progress_bar(pct=pct_stage)} preparing links {link_call_idx}/{total_link_calls}"
                    )
                payload = self._request(
                    "GET",
                    ENTRY_DOWNLOAD_LINKS_PATH,
                    params={
                        "record_id": int(record_id),
                        "symbol": str(sym),
                        "start_timestamp": int(s_us),
                        "end_timestamp": int(e_us),
                    },
                    auth="api_key",
                )
                items0 = _normalize_download_links(payload)
                # Tag each item with its symbol for later row labeling.
                for it in items0:
                    if isinstance(it, dict):
                        it["_symbol"] = sym
                all_items.extend(items0)

        total_files = len([x for x in all_items if isinstance(x, dict) and x.get("download_url")])
        done_files = 0
        total_bytes = 0
        downloaded_total = 0

        # Best-effort: sum Content-Length for a truer overall percentage.
        # If unavailable, we fallback to file-count-based percentage.
        urls0 = [x.get("download_url") for x in all_items if isinstance(x, dict) and x.get("download_url")]
        for idx, url0 in enumerate(urls0, 1):
            if progress.enabled:
                pct_stage = (idx / max(len(urls0), 1)) * 100.0
                progress.update(
                    f"Downloading {_format_progress_bar(pct=pct_stage)} probing sizes {idx}/{len(urls0)}"
                )
            if not url0:
                continue
            try:
                h = session.head(str(url0), timeout=dl_timeout, allow_redirects=True)
                if h.status_code < 400:
                    cl = h.headers.get("Content-Length")
                    if cl:
                        total_bytes += int(cl)
            except Exception:
                continue

        last_render = 0.0

        for item in all_items:
            download_url = item.get("download_url") if isinstance(item, dict) else None
            path = item.get("path") if isinstance(item, dict) else None
            sym_name = item.get("_symbol") if isinstance(item, dict) else None
            if not download_url:
                continue

            # Download signed URL with a small retry loop (network/proxy can reset connections).
            last_dl_err: Optional[Exception] = None
            resp = None
            for attempt in range(5):
                try:
                    resp = session.get(str(download_url), timeout=dl_timeout, stream=True)
                    last_dl_err = None
                    break
                except requests.RequestException as e:
                    last_dl_err = e
                    time.sleep(0.5 * (2**attempt))
            if resp is None:
                raise RequestError(f"Failed to download file after retries: path={path} err={last_dl_err}")
            if resp.status_code >= 400:
                raise RequestError(f"Failed to download file: HTTP {resp.status_code} path={path}")

            file_total = 0
            try:
                file_total = int(resp.headers.get("Content-Length") or 0)
            except Exception:
                file_total = 0

            file_downloaded = 0
            buf = bytearray()

            for chunk in resp.iter_content(chunk_size=256 * 1024):
                if not chunk:
                    continue
                buf += chunk
                file_downloaded += len(chunk)
                now = time.time()
                if now - last_render >= 0.15:
                    # Overall fraction: done files + fraction of current file (if we know its length)
                    frac_current = (file_downloaded / max(file_total, 1)) if file_total > 0 else 0.0
                    overall_files_frac = (done_files + frac_current) / max(total_files, 1)
                    pct = overall_files_frac * 100.0

                    downloaded_total_now = downloaded_total + file_downloaded
                    if total_bytes > 0:
                        # Prefer byte-based percent when possible.
                        pct_bytes = (downloaded_total_now / max(total_bytes, 1)) * 100.0
                        pct = pct_bytes
                    progress.update(
                        f"Downloading {_format_progress_bar(pct=pct)} "
                        f"{done_files}/{total_files} {_format_bytes(downloaded_total_now)}"
                        + (f"/{_format_bytes(total_bytes)}" if total_bytes > 0 else "")
                    )
                    last_render = now

            downloaded_total += file_downloaded
            done_files += 1

            # Parse protobuf messages
            messages = parse_length_prefixed_messages(bytes(buf), msg_cls)
            decrypt_key = _normalize_encryption_key(
                encryption_key
                or os.getenv("LIBALPHA_RECORD_ENCRYPTION_KEY")
            )
            for msg in messages:
                symbol_values = list(getattr(msg, "symbol_values", []) or [])

                for sv in symbol_values:
                    sid = int(getattr(sv, "symbol_id", 0) or 0)
                    if sid <= 0:
                        continue
                    payload_kind = None
                    try:
                        payload_kind = sv.WhichOneof("payload")
                    except Exception:
                        payload_kind = None

                    features_list: List[float] = []
                    if payload_kind == "values":
                        dv = getattr(sv, "values", None)
                        items = getattr(dv, "items", None) if dv is not None else None
                        features_list = list(items) if items is not None else []
                    elif payload_kind == "encrypted_payload":
                        enc = getattr(sv, "encrypted_payload", None)
                        if enc and len(enc) > 0 and decrypt_key:
                            decrypted = decrypt_data_entry_payload(decrypt_key, bytes(enc))
                            if decrypted is not None:
                                features_list = decrypted

                    d = {
                        "entry_id": int(getattr(msg, "entry_id", 0)),
                        "record_id": int(getattr(msg, "record_id", 0)),
                        "symbol": str(sym_name) if sym_name else "",
                        "symbol_id": int(sid),
                        "features": features_list,
                        "timestamp": int(getattr(msg, "timestamp", 0)),
                        "runner_timestamp": int(getattr(msg, "runner_timestamp", 0)),
                        "server_timestamp": int(getattr(msg, "server_timestamp", 0)),
                        "request_id": str(getattr(msg, "request_id", "") or ""),
                    }
                    as_ts = getattr(msg, "as_of_timestamp", None)
                    if as_ts is not None:
                        d["as_of_timestamp"] = int(as_ts)
                    if path:
                        d["source_path"] = str(path)
                    rows.append(d)

        # Final overall 100% render
        if progress.enabled:
            progress.update(
                f"Downloading {_format_progress_bar(pct=100.0)} {done_files}/{total_files} "
                f"{_format_bytes(downloaded_total)}" + (f"/{_format_bytes(total_bytes)}" if total_bytes > 0 else "")
            )
        progress.finish()
        return pd.DataFrame(rows)

    # ============================================================
    # Historical Upload API (Public): upload_data(record_id, df, ...)->None
    # ============================================================
    def upload_data(
        self,
        record_id: int,
        df: "pd.DataFrame",
        encryption_key: Optional[str] = None,
    ) -> None:
        """
        Historical backfill uploader (protobuf length-prefixed DataEntry stream).

        This method delegates to `liberal_alpha.history_upload.upload_data` and exists
        primarily as a convenience wrapper so callers can use one `client` object.

        Args:
            record_id: Backend record identifier.
            df: DataFrame accepted by `history_upload.upload_data` (requires at least
                `symbol` and `timestamp`, plus either `features` or feature columns).
            encryption_key: Record encryption key for encrypting payload when the record is
                encrypted. If omitted, env `LIBALPHA_RECORD_ENCRYPTION_KEY` is used.
        """
        from .history_upload import upload_data as _upload_data

        if not self.api_key:
            raise ConfigurationError("Missing api_key (LIBALPHA_API_KEY). upload_data needs auth.")
        api_key_eff = self.api_key
        record_pk = encryption_key or os.getenv("LIBALPHA_RECORD_ENCRYPTION_KEY")

        _upload_data(record_id=int(record_id), df=df, api_key=api_key_eff, encryption_key=record_pk)

    # ============================================================
    # ✅ V2 additions: gRPC send_data() (Runner) - does not affect existing HTTP logic
    # ============================================================
    def _grpc_pick_stub_and_request(self):
        """
        Pick gRPC Stub + Request message class from generated proto module.

        Compatible with your current generated file:
        - service_pb2_grpc.JsonServiceStub
        - service_pb2.JsonRequest
        Also works if names change (falls back to any *Stub / *Request).
        """
        try:
            import grpc  # noqa: F401
        except Exception as e:
            raise ConfigurationError("grpcio is required for send_data(). Please `pip install grpcio`.") from e

        try:
            from .proto import service_pb2_grpc, service_pb2
        except Exception as e:
            raise ConfigurationError(
                "Cannot import proto modules. Expected liberal_alpha/proto/service_pb2.py and service_pb2_grpc.py"
            ) from e

        # ---- pick stub ----
        stub_cls = None
        if hasattr(service_pb2_grpc, "JsonServiceStub"):
            stub_cls = getattr(service_pb2_grpc, "JsonServiceStub")
        else:
            for name in dir(service_pb2_grpc):
                if name.endswith("Stub"):
                    stub_cls = getattr(service_pb2_grpc, name)
                    break
        if stub_cls is None:
            raise ConfigurationError("Cannot find gRPC stub class in proto module (service_pb2_grpc).")

        # ---- pick request msg ----
        req_cls = None
        if hasattr(service_pb2, "JsonRequest"):
            req_cls = getattr(service_pb2, "JsonRequest")
        else:
            for name in dir(service_pb2):
                if name.endswith("Request"):
                    obj = getattr(service_pb2, name)
                    if hasattr(obj, "DESCRIPTOR"):
                        req_cls = obj
                        break
        if req_cls is None:
            raise ConfigurationError("Cannot find gRPC request message class in proto module (service_pb2).")

        return stub_cls, req_cls

    def _grpc_build_request(self, req_cls, identifier: str, data: dict, record_id: int, event_type: str):
        """
        Build protobuf request message in a field-safe way:
        - Only set fields that actually exist in the generated message schema.

        Your current proto(JsonRequest) is typically:
          - json_data: string
          - event_type: string
          - timestamp: int64
          - metadata: map<string,string>
        """
        # discover fields
        try:
            field_names = {f.name for f in req_cls.DESCRIPTOR.fields}
        except Exception:
            field_names = set()

        ts_ms = int(time.time() * 1000)

        # Put everything (including record_id/identifier) into json_data to be proto-compatible
        payload_json = json.dumps(data or {}, ensure_ascii=False)


        md = {
            "record_id": str(record_id),
            "identifier": str(identifier),
            "event_type": str(event_type),
        }

        kwargs: Dict[str, Any] = {}

        # common request shapes
        if "json_data" in field_names:
            kwargs["json_data"] = payload_json
        elif "jsonData" in field_names:
            kwargs["jsonData"] = payload_json
        elif "data" in field_names:
            # could be string or bytes; try string first
            kwargs["data"] = payload_json
        elif "payload" in field_names:
            kwargs["payload"] = payload_json
        elif "message" in field_names:
            kwargs["message"] = payload_json

        if "event_type" in field_names:
            kwargs["event_type"] = str(event_type)
        elif "eventType" in field_names:
            kwargs["eventType"] = str(event_type)

        if "timestamp" in field_names:
            kwargs["timestamp"] = ts_ms
        elif "timestamp_ms" in field_names:
            kwargs["timestamp_ms"] = ts_ms
        elif "timestampMs" in field_names:
            kwargs["timestampMs"] = ts_ms

        if "metadata" in field_names:
            kwargs["metadata"] = md

        # create message
        try:
            req = req_cls(**kwargs)
        except TypeError:
            req = req_cls()
            for k, v in kwargs.items():
                try:
                    setattr(req, k, v)
                except Exception:
                    pass

        # If schema expects bytes for 'data', convert
        try:
            if "data" in field_names:
                f = req_cls.DESCRIPTOR.fields_by_name.get("data")
                # TYPE_BYTES == 12
                if f is not None and getattr(f, "type", None) == 12 and isinstance(getattr(req, "data", None), str):
                    setattr(req, "data", payload_json.encode("utf-8"))
        except Exception:
            pass

        return req

    def send_data(
        self,
        *,
        identifier: str,
        data: dict,
        record_id: int,
        event_type: str = "raw",
        timeout: Optional[float] = None,
    ):
        """
        Send one JSON payload to local gRPC runner.

        Compatible with generated service:
        - stub: JsonServiceStub
        - rpc: ProcessJson
        - request: JsonRequest (field-safe population)
        """
        # ---- runner host/port ----
        host = (
            getattr(self, "host", None)
            or getattr(self, "runner_host", None)
            or os.getenv("LIBALPHA_RUNNER_HOST", "127.0.0.1")
        )
        port = int(
            getattr(self, "port", None)
            or getattr(self, "runner_port", None)
            or os.getenv("LIBALPHA_RUNNER_PORT", "8128")
        )
        target = f"{host}:{port}"

        # ---- grpc deps ----
        import grpc

        stub_cls, req_cls = self._grpc_pick_stub_and_request()

        req = self._grpc_build_request(
            req_cls=req_cls,
            identifier=str(identifier),
            data=data or {},
            record_id=int(record_id),
            event_type=str(event_type),
        )

        # ---- call ----
        call_timeout = float(timeout) if timeout is not None else float(getattr(self, "timeout", 30))

        try:
            with grpc.insecure_channel(target) as channel:
                stub = stub_cls(channel)

                if hasattr(stub, "ProcessJson"):
                    return stub.ProcessJson(req, timeout=call_timeout)

                # fallback: find a callable attr
                for name in dir(stub):
                    if name.startswith("_"):
                        continue
                    fn = getattr(stub, name, None)
                    if callable(fn):
                        return fn(req, timeout=call_timeout)

                raise ConfigurationError("No callable RPC method found on gRPC stub.")
        except grpc.RpcError as e:
            raise RequestError(f"gRPC call failed: {e}") from e

    # ============================================================
    # subscribe_data: 实时订阅 ws/data，按 symbol 过滤，回调 on_data_fn(symbol, DataFrame)
    # ============================================================
    def subscribe_data(
        self,
        record_id: int,
        *,
        symbols: List[str],
        on_data_fn: Callable[[str, "pd.DataFrame"], None],
        encryption_key: Optional[str] = None,
    ) -> None:
        """Subscribes to live data for a record into a DataFrame (schema matches publish).

        Args:
            record_id: Backend record identifier.
            symbols: Symbols to filter by (required).
            on_data_fn: Callback function invoked upon receiving new data. The
                function takes two arguments: the symbol (str) and a DataFrame
                containing the new data for that symbol.
            encryption_key: Optional, record encryption key for decrypting encrypted_payload.
                Required if data is encrypted (encrypted_payload has value).
        """
        _ensure_pandas()
        if not self.api_key:
            raise ConfigurationError("subscribe_data requires api_key (LIBALPHA_API_KEY).")
        if not symbols:
            raise ValueError("symbols is required and cannot be empty.")

        payload = self.my_subscriptions()
        data = payload.get("data") or {}
        subs = data.get("subscriptions") or []
        record_ids = []
        for s in subs:
            rid = s.get("record_id") or (s.get("record") or {}).get("record_id")
            if rid is not None:
                record_ids.append(int(rid))
        if record_id not in record_ids:
            print(
                f"record_id={int(record_id)} is not accessible. "
                "You may not be subscribed to this record."
            )
            return

        # subscribe_data accepts symbol NAME strings; ws payload uses symbol_id (stringified).
        # Map name -> symbol_id using /api/records/{record_id} target_symbols so subscribed
        # records owned by other users can still be resolved.
        symbols_in = [str(s).strip() for s in (symbols or []) if str(s).strip()]
        if not symbols_in:
            raise ValueError("symbols is required and cannot be empty.")

        need_map = any((not s.isdigit()) for s in symbols_in)
        name_to_id: Dict[str, str] = {}
        if need_map:
            try:
                rec = self._request("GET", f"/api/records/{int(record_id)}", auth="api_key")
            except RequestError as e:
                msg = str(e)
                if "HTTP 403" in msg or "HTTP 404" in msg:
                    print(
                        f"record_id={int(record_id)} is not accessible. "
                        "You may not be subscribed to this record."
                    )
                    return
                raise
            ts = (rec.get("data") or {}).get("target_symbols") if isinstance(rec, dict) else None
            if not isinstance(ts, list):
                ts = []
            for item in ts:
                if not isinstance(item, dict):
                    continue
                sym = item.get("symbol")
                sid = item.get("symbol_id")
                if isinstance(sym, str) and sym.strip() and sid is not None:
                    try:
                        sid_int = int(sid)
                    except Exception:
                        continue
                    if sid_int > 0:
                        name_to_id[sym.strip()] = str(sid_int)
            if not name_to_id:
                raise ConfigurationError(
                    f"Cannot map symbol names to symbol_id from /api/records/{int(record_id)}. "
                    "You may not be subscribed to this record, the record may not expose target_symbols, "
                    "or you can pass symbol_id strings directly."
                )

        # Build symbol_id filters (as strings) + mapping back to original input symbol string.
        symbol_ids: List[str] = []
        id_to_input: Dict[str, str] = {}
        for s in symbols_in:
            if s.isdigit():
                sid = s
            else:
                sid = name_to_id.get(s)
                if not sid:
                    raise ConfigurationError(f"Unknown symbol (no symbol_id mapping): {s!r}")
            symbol_ids.append(sid)
            # Prefer first occurrence; stable mapping to original user input
            if sid not in id_to_input:
                id_to_input[sid] = s

        record_pk = (
            encryption_key
            or os.getenv("LIBALPHA_RECORD_ENCRYPTION_KEY")
        )
        decrypt_key = _normalize_encryption_key(record_pk)
        max_reconnect = 3
        wallet = get_user_wallet_address(self.api_base, self.api_key or "")

        try:
            from .proto import data_entry_pb2
        except Exception as e:
            raise ConfigurationError("Cannot import data_entry_pb2 for subscribe_data.") from e

        ConsumerMessage = data_entry_pb2.ConsumerMessage
        ws_url = f"{_to_ws_url(self.api_base)}/ws/data"
        symbol_set = set(symbol_ids)
        rec_id = int(record_id)
        stop_requested = False

        def _entry_to_rows(entry: Any) -> List[Dict[str, Any]]:
            symbol_values = list(getattr(entry, "symbol_values", []) or [])

            base = {
                "entry_id": int(getattr(entry, "entry_id", 0)),
                "record_id": int(getattr(entry, "record_id", 0)),
                "timestamp": int(getattr(entry, "timestamp", 0)),
                "runner_timestamp": int(getattr(entry, "runner_timestamp", 0)),
                "server_timestamp": int(getattr(entry, "server_timestamp", 0)),
                "request_id": str(getattr(entry, "request_id", "") or ""),
            }
            as_ts = getattr(entry, "as_of_timestamp", None)
            if as_ts is not None:
                base["as_of_timestamp"] = int(as_ts)

            out: List[Dict[str, Any]] = []
            for sv in symbol_values:
                sid = int(getattr(sv, "symbol_id", 0) or 0)
                if sid <= 0:
                    continue

                payload_kind = None
                try:
                    payload_kind = sv.WhichOneof("payload")
                except Exception:
                    payload_kind = None

                features_list: List[float] = []
                if payload_kind == "values":
                    dv = getattr(sv, "values", None)
                    items = getattr(dv, "items", None) if dv is not None else None
                    features_list = list(items) if items is not None else []
                elif payload_kind == "encrypted_payload":
                    enc = getattr(sv, "encrypted_payload", None)
                    if enc and len(enc) > 0:
                        if not decrypt_key:
                            raise ConfigurationError(
                                "Encrypted data received. Please provide encryption_key to decrypt this record."
                            )
                        decrypted = decrypt_data_entry_payload(decrypt_key, bytes(enc))
                        if decrypted is None:
                            raise ConfigurationError(
                                "Failed to decrypt encrypted data. The provided encryption_key may be incorrect."
                            )
                        features_list = decrypted

                out.append(dict(base, symbol=str(sid), features=features_list))
            return out

        def _handle_json_data(msg: Dict[str, Any], rec: int) -> Optional[Dict[str, Any]]:
            if msg.get("status") != "data" or "data" not in msg:
                return None
            raw = msg["data"]
            if isinstance(raw, str):
                try:
                    raw = base64.b64decode(raw)
                except Exception:
                    raw = raw.encode("utf-8") if isinstance(raw, str) else raw
            if isinstance(raw, str):
                return None
            payload_bytes = bytes(raw) if not isinstance(raw, bytes) else raw
            if len(payload_bytes) < 12:
                return None
            if not decrypt_key:
                raise ConfigurationError("Encrypted data received. Please provide encryption_key to decrypt this record.")
            decrypted = decrypt_data_entry_payload(decrypt_key, payload_bytes)
            if decrypted is None:
                raise ConfigurationError(
                    "Failed to decrypt encrypted data. The provided encryption_key may be incorrect."
                )
            ts = msg.get("timestamp") or 0
            return {
                "entry_id": 0,
                "record_id": rec,
                "symbol": "",
                "features": decrypted,
                "timestamp": ts,
                "runner_timestamp": ts,
                "server_timestamp": ts,
                "request_id": "",
                "batch_id": 0,
            }

        try:
            import websocket  # type: ignore
        except Exception as e:
            raise ConfigurationError("websocket-client not installed. `pip install websocket-client`") from e

        def _on_message(ws: Any, msg: Any) -> None:
            nonlocal stop_requested
            try:
                rows_out: List[Dict[str, Any]] = []
                if isinstance(msg, bytes):
                    try:
                        cm = ConsumerMessage()
                        cm.ParseFromString(msg)
                        entry = cm.entry
                        if entry is not None:
                            rows_out = _entry_to_rows(entry)
                    except Exception:
                        pass
                else:
                    try:
                        obj = json.loads(msg) if isinstance(msg, str) else msg
                        if isinstance(obj, dict):
                            row = _handle_json_data(obj, rec_id)
                            if row is not None:
                                rows_out = [row]
                    except Exception:
                        pass
                for row in rows_out:
                    sym = row.get("symbol") or ""
                    if sym in symbol_set:
                        on_data_fn(id_to_input.get(sym, sym), pd.DataFrame([row]))
                    else:
                        logger.debug(
                            "subscribe_data: symbol '%s' not in filter %s, skip",
                            sym, list(symbol_set),
                        )
            except ConfigurationError as e:
                print(str(e))
                stop_requested = True
                try:
                    ws.close()
                except Exception:
                    pass
            except Exception as e:
                logger.exception("subscribe_data on_message error: %s", e)

        def _on_open(ws: Any) -> None:
            ws.send(json.dumps({"wallet_address": wallet, "record_id": rec_id}))

        attempts = 0
        last_err: Optional[Exception] = None
        while attempts <= int(max_reconnect):
            attempts += 1
            try:
                wsapp = websocket.WebSocketApp(ws_url, on_message=_on_message, on_open=_on_open)
                wsapp.run_forever(ping_interval=20, ping_timeout=10)
            except Exception as e:
                last_err = e
            if stop_requested:
                return
            if attempts <= int(max_reconnect):
                time.sleep(min(2 ** (attempts - 1), 10))
        raise RequestError(f"WebSocket subscribe_data failed after retries. last_err={last_err}")


# ----------------------------
# Keep compatibility with liberal_alpha/__init__.py
# ----------------------------
liberal: Optional[LiberalAlphaClient] = None


def initialize(
    api_key: Optional[str] = None,
    api_base: str = DEFAULT_API_BASE,
    timeout: int = 30,
) -> LiberalAlphaClient:
    global liberal
    liberal = LiberalAlphaClient(
        api_key=api_key,
        api_base=api_base,
        timeout=timeout,
    )
    return liberal
