from __future__ import annotations

import hashlib
import math
import os
import struct
import time
import sys
import shutil
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

import requests

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

from Crypto.Cipher import AES

from .exceptions import ConfigurationError, RequestError
from .proto import data_entry_pb2


API_BASE_FIXED = "https://api.librealpha.com"
DEFAULT_TIMEOUT_S = 20
HISTORY_UPLOAD_PATH = "/api/entries/history/upload"

_HTTP_SESSION: Optional[requests.Session] = None


def _get_http_session() -> requests.Session:
    """
    Reuse a single requests.Session and respect LIBALPHA_HTTP_TRUST_ENV.
    In some Windows/corp proxy setups, disabling trust_env can avoid TLS EOF issues.
    """
    global _HTTP_SESSION
    if _HTTP_SESSION is None:
        s = requests.Session()
        s.trust_env = os.getenv("LIBALPHA_HTTP_TRUST_ENV", "1") != "0"
        _HTTP_SESSION = s
    return _HTTP_SESSION


_SYSTEM_COLS = {
    "symbol",
    "timestamp",
    "batch_id",
    "as_of_time",
    "as_of_timestamp",
    "request_id",
    "features",  # optional shortcut
}


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


def _is_nan(x: Any) -> bool:
    try:
        return x is None or (isinstance(x, float) and math.isnan(x)) or bool(pd.isna(x))  # type: ignore[attr-defined]
    except Exception:
        return x is None


def _guess_unix_unit_from_int(v: int) -> str:
    """
    Heuristic: seconds vs ms vs us (for converting integer timestamps).
    """
    av = abs(int(v))
    if av < 100_000_000_000:  # < 1e11 => seconds
        return "s"
    if av < 100_000_000_000_000:  # < 1e14 => milliseconds
        return "ms"
    return "us"


def _to_datetime_utc(series) -> "pd.Series":
    _ensure_pandas()
    # datetime dtype
    if pd.api.types.is_datetime64_any_dtype(series):  # type: ignore[union-attr]
        out = pd.to_datetime(series, utc=True, errors="coerce")  # type: ignore[union-attr]
        if out.isna().any():
            raise ValueError("timestamp column contains invalid datetimes")
        return out

    # numeric dtype
    if pd.api.types.is_numeric_dtype(series):  # type: ignore[union-attr]
        s = series.dropna()
        if len(s) == 0:
            raise ValueError("timestamp column is empty")
        unit = _guess_unix_unit_from_int(int(s.iloc[0]))
        out = pd.to_datetime(series.astype("int64"), unit=unit, utc=True, errors="coerce")  # type: ignore[union-attr]
        if out.isna().any():
            raise ValueError("timestamp column contains invalid numeric timestamps")
        return out

    # string / object
    out = pd.to_datetime(series, utc=True, errors="coerce")  # type: ignore[union-attr]
    if out.isna().any():
        raise ValueError("timestamp column contains invalid timestamps (expected datetime64 or unix ts)")
    return out


def _datetime_series_to_us(ts_utc: "pd.Series") -> "pd.Series":
    """
    Convert pandas datetime series to int microseconds since epoch (UTC).
    """
    # Robust across pandas versions: convert to numpy datetime64[ns] then to int64 ns.
    arr = ts_utc.to_numpy(dtype="datetime64[ns]", copy=False)
    ns = arr.astype("int64")
    return pd.Series(ns // 1000, index=ts_utc.index, dtype="int64")


def _check_dataframe(df: "pd.DataFrame") -> "pd.DataFrame":
    """
    Normalize and validate input DataFrame.

    Required columns:
      - symbol (str)
      - timestamp (datetime64 or unix timestamp; normalized to datetime64[ns, UTC])

    Optional system columns:
      - batch_id (int)
      - as_of_time / as_of_timestamp (datetime or unix ts)
      - request_id (str)
      - features (list[float] or string like "[1.0, nan, 2.0]")
    """
    _ensure_pandas()
    if df is None or getattr(df, "empty", True):
        raise ValueError("df is empty")

    if "symbol" not in df.columns:
        raise ValueError("DataFrame missing required column: symbol")
    if "timestamp" not in df.columns:
        raise ValueError("DataFrame missing required column: timestamp")

    out = df.copy()
    out["symbol"] = out["symbol"].astype(str).str.strip()
    if (out["symbol"] == "").any():
        raise ValueError("symbol column contains empty values")

    out["timestamp"] = _to_datetime_utc(out["timestamp"])

    # Normalize as_of_time/as_of_timestamp if present
    if "as_of_time" in out.columns and "as_of_timestamp" not in out.columns:
        out["as_of_timestamp"] = _to_datetime_utc(out["as_of_time"])
    if "as_of_timestamp" in out.columns:
        out["as_of_timestamp"] = _to_datetime_utc(out["as_of_timestamp"])

    return out


def _parse_feature_list_str(s: str) -> List[float]:
    """
    Parse a string like "[1.0, nan, 2.0]" safely (no eval).
    Accepts "nan"/"NaN" tokens.
    """
    s = (s or "").strip()
    if not s:
        return []
    if s.startswith("[") and s.endswith("]"):
        s = s[1:-1]
    s = s.strip()
    if not s:
        return []
    parts = s.split(",")
    out: List[float] = []
    for p in parts:
        t = p.strip()
        if not t:
            continue
        if t.lower() == "nan":
            out.append(float("nan"))
            continue
        try:
            out.append(float(t))
        except Exception:
            # tolerate "None"/"null"
            if t.lower() in ("none", "null"):
                out.append(float("nan"))
                continue
            raise ValueError(f"Invalid feature token: {t!r}")
    return out


def _encrypt_features_aes256gcm_record_key(record_private_key: str, features: Sequence[float]) -> bytes:
    """
    Compatible with liberal_runner/internal/memory/encryption_v2.go:
      key = sha256(record_private_key)
      plaintext = uint32_le(count) || float64_le[count]
      output = nonce(12) || ciphertext||tag
    """
    key = hashlib.sha256(record_private_key.encode("utf-8")).digest()
    plain = struct.pack("<I", int(len(features)))
    for v in features:
        fv = float(v)
        plain += struct.pack("<d", fv)
    nonce = os.urandom(12)
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    ciphertext, tag = cipher.encrypt_and_digest(plain)
    return nonce + ciphertext + tag


def _length_prefix_be(n: int) -> bytes:
    if n < 0:
        raise ValueError("length must be non-negative")
    return int(n).to_bytes(4, byteorder="big", signed=False)


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
        self._stream = sys.stderr

    def update(self, text: str) -> None:
        if not self.enabled:
            return
        s = str(text).replace("\r", " ").replace("\n", " ")

        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))
        self._stream.write("\r" + s + pad)
        self._stream.flush()
        self._last_len = len(s)

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


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 _request_with_retries(
    method: str,
    url: str,
    *,
    headers: Dict[str, str],
    params: Optional[Dict[str, Any]] = None,
    data: Optional[bytes] = None,
    timeout: int = DEFAULT_TIMEOUT_S,
    retries: int = 5,
) -> requests.Response:
    """
    Minimal retry wrapper for transient network/SSL issues on Windows.
    """
    last_err: Optional[Exception] = None
    for i in range(max(int(retries), 1)):
        try:
            session = _get_http_session()
            return session.request(method, url, headers=headers, params=params, data=data, timeout=int(timeout))
        except requests.RequestException as e:
            last_err = e
            # Exponential-ish backoff: 0.5s, 1s, 2s...
            time.sleep(0.5 * (2**i))
    raise RequestError(f"HTTP {method} {url} failed after retries: {last_err}")


def _get_json_with_retries(
    url: str,
    *,
    headers: Dict[str, str],
    timeout: int = DEFAULT_TIMEOUT_S,
    retries: int = 5,
) -> Tuple[int, Any]:
    resp = _request_with_retries("GET", url, headers=headers, timeout=timeout, retries=retries)
    try:
        return resp.status_code, resp.json()
    except Exception:
        return resp.status_code, resp.text


@dataclass(frozen=True)
class _MinuteKey:
    symbol: str
    minute_utc: int  # unix minute bucket (timestamp_us // 60_000_000)


def _fetch_record_feature_names(*, api_base: str, api_key: str, record_id: int, timeout: int = 30) -> List[str]:
    """
    Uses /api/records/{id} and reads data.features_original[].name
    """
    url = f"{api_base}/api/records/{int(record_id)}"
    status, payload = _get_json_with_retries(
        url,
        headers={"X-API-Key": api_key, "Accept": "application/json"},
        timeout=int(timeout),
        retries=3,
    )
    if int(status) >= 400:
        raise RequestError(f"Failed to fetch record meta: HTTP {status} {str(payload).strip()[:300]}")
    if not isinstance(payload, dict):
        raise RequestError(f"Record meta response not JSON: {str(payload).strip()[:200]}")

    data = payload.get("data") if isinstance(payload, dict) else None
    if not isinstance(data, dict):
        return []
    feats = data.get("features_original") or data.get("features") or []
    if not isinstance(feats, list):
        return []
    out: List[str] = []
    for item in feats:
        if isinstance(item, dict):
            name = item.get("name")
            if isinstance(name, str) and name.strip():
                out.append(name.strip())
        elif isinstance(item, str) and item.strip():
            out.append(item.strip())
    return out


def _record_is_encrypted(*, api_base: str, api_key: str, record_id: int, timeout: int = 30) -> Optional[bool]:
    """
    Best-effort: look for encryption flags in /api/records/{id}.
    Returns:
      - True/False if a flag is found
      - None if unknown (API doesn't expose it)
    """
    url = f"{api_base}/api/records/{int(record_id)}"
    try:
        status, payload = _get_json_with_retries(
            url,
            headers={"X-API-Key": api_key, "Accept": "application/json"},
            timeout=int(timeout),
            retries=5,
        )
        if int(status) < 400:
            data = payload.get("data") if isinstance(payload, dict) else None
            if isinstance(data, dict):
                for k in ("is_encrypted", "isEncrypted", "encrypted", "encryption_enabled", "encryptionEnabled"):
                    v = data.get(k)
                    if isinstance(v, bool):
                        return v
    except Exception:
        pass

    # Fallback: /api/records list includes is_encrypted (but /api/records/{id} currently omits it)
    try:
        url2 = f"{api_base}/api/records"
        status2, payload2 = _get_json_with_retries(
            url2,
            headers={"X-API-Key": api_key, "Accept": "application/json"},
            timeout=int(timeout),
            retries=5,
        )
        if int(status2) >= 400 or not isinstance(payload2, dict):
            return None
        data2 = payload2.get("data") if isinstance(payload2, dict) else None
        if not isinstance(data2, dict):
            return None
        recs = data2.get("records")
        if not isinstance(recs, list):
            return None
        rid = int(record_id)
        for r in recs:
            if not isinstance(r, dict):
                continue
            r_id = r.get("id", None)
            if r_id is None:
                r_id = r.get("record_id", None)
            try:
                if int(r_id) != rid:
                    continue
            except Exception:
                continue
            v = r.get("is_encrypted")
            if isinstance(v, bool):
                return v
            return None
    except Exception:
        return None


def _fetch_record_symbol_id_map(*, api_base: str, api_key: str, record_id: int, timeout: int = 30) -> Dict[str, int]:
    """
    Fetch target symbol mapping for a record.

    Uses /api/records/user-records and reads data.records[].target_symbols[{symbol_id, symbol}].
    Returns:
      dict: symbol_name -> symbol_id
    """
    url = f"{api_base}/api/records/user-records"
    status, payload = _get_json_with_retries(
        url,
        headers={"X-API-Key": api_key, "Accept": "application/json"},
        timeout=int(timeout),
        retries=3,
    )
    if int(status) >= 400:
        raise RequestError(f"Failed to fetch user records: HTTP {status} {str(payload).strip()[:300]}")
    if not isinstance(payload, dict):
        raise RequestError(f"User records response not JSON: {str(payload).strip()[:200]}")

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

    rid = int(record_id)
    for r in recs:
        if not isinstance(r, dict):
            continue
        try:
            if int(r.get("id")) != rid:
                continue
        except Exception:
            continue
        ts = r.get("target_symbols")
        if not isinstance(ts, list):
            return {}
        out: Dict[str, int] = {}
        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():
                try:
                    out[sym.strip()] = int(sid)
                except Exception:
                    continue
        return out

    return {}


def upload_data(
    record_id: int,
    df: "pd.DataFrame",
    api_key: Optional[str] = None,
    encryption_key: Optional[str] = None,
) -> None:
    """Uploads a batch of points for a given record via DataFrame.

    Args:
        record_id: Backend record identifier to attach the rows to.
        df: DataFrame with required columns ``symbol`` (str), ``timestamp``
            (datetime64), and optionally opt-in system columns like batch_id and
            as_of_time, and any number of user-defined float columns matching the record
            schema. User columns must be convertible to double and may contain NaN.
        api_key: Optional. API key used as X-API-Key. Falls back to env LIBALPHA_API_KEY.
        encryption_key: Optional. Record encryption key (string) used to AES-GCM encrypt features into
            DataEntry.encrypted_payload when the record is encrypted.
        Note: API base is fixed to https://api.librealpha.com (same as C++ uploader).
    """
    _ensure_pandas()
    if record_id <= 0:
        raise ValueError("record_id must be positive int")

    api_base = API_BASE_FIXED
    api_key = (api_key or os.getenv("LIBALPHA_API_KEY") or "").strip()
    if not api_key:
        raise ConfigurationError("Missing api_key (LIBALPHA_API_KEY).")

    df = _check_dataframe(df)

    # Determine feature order if not passing pre-packed "features"
    use_features_col = "features" in df.columns
    feature_names: List[str] = []
    if not use_features_col:
        feature_names = _fetch_record_feature_names(
            api_base=api_base, api_key=api_key, record_id=int(record_id), timeout=int(DEFAULT_TIMEOUT_S)
        )
        if not feature_names:
            raise ConfigurationError(
                "Cannot determine record feature order from /api/records/{id}. "
                "Either provide a 'features' column, or ensure the backend returns features_original."
            )

    # Determine whether to encrypt (best effort)
    enc_flag = _record_is_encrypted(api_base=api_base, api_key=api_key, record_id=int(record_id), timeout=int(DEFAULT_TIMEOUT_S))
    encrypt = bool(enc_flag)  # True only when backend explicitly says encrypted
    if encrypt and not (encryption_key or "").strip():
        raise ConfigurationError("Record is encrypted but encryption_key was not provided.")

    # Prepare timestamp microseconds + grouping key
    ts_us = _datetime_series_to_us(df["timestamp"])
    symbols_raw = df["symbol"].astype(str).apply(lambda x: str(x).strip())

    # Resolve symbol -> symbol_id string key (new proto uses symbol_id).
    sym_id_map: Dict[str, int] = {}
    if any((s and not s.isdigit()) for s in symbols_raw.tolist()):
        sym_id_map = _fetch_record_symbol_id_map(
            api_base=api_base, api_key=api_key, record_id=int(record_id), timeout=int(DEFAULT_TIMEOUT_S)
        )
        if not sym_id_map:
            raise ConfigurationError(
                "DataFrame 'symbol' contains non-numeric values, but backend did not return target_symbols mapping. "
                "Please pass symbol_id strings (e.g. '1','3') in df['symbol']."
            )

    def _to_symbol_key(s: str) -> str:
        s2 = (s or "").strip()
        if not s2:
            return ""
        if s2.isdigit():
            return s2
        sid = sym_id_map.get(s2, 0)
        if not sid:
            raise ConfigurationError(f"Unknown symbol (no symbol_id mapping): {s2!r}")
        return str(int(sid))

    symbols = symbols_raw.apply(_to_symbol_key)
    minute_bucket = (ts_us // 60_000_000).astype("int64")

    # Build per-minute indices (we pack multiple symbol_id into symbol_values).
    groups: Dict[int, List[int]] = {}
    for i, mb in enumerate(minute_bucket.tolist()):
        groups.setdefault(int(mb), []).append(i)

    progress = _SingleLineProgress(enabled=os.getenv("LIBALPHA_UPLOAD_PROGRESS", "1") != "0")
    total_groups = len(groups)
    done_groups = 0
    sent_bytes = 0

    headers = {
        "X-API-Key": api_key,
        "Content-Type": "application/octet-stream",
        "Accept": "application/json",
    }

    # Process in sorted order for determinism
    keys_sorted = sorted(groups.keys())
    try:
        for mb in keys_sorted:
            idxs = groups[mb]

            # Enforce "each request = 1 minute of data" (minute bucket is the request unit)
            buf = bytearray()
            # Within the minute, aggregate rows by exact timestamp_us.
            by_ts: Dict[int, List[int]] = {}
            for i in idxs:
                by_ts.setdefault(int(ts_us.iloc[i]), []).append(i)

            for t_us, t_idxs in sorted(by_ts.items(), key=lambda x: x[0]):
                entry = data_entry_pb2.DataEntry()
                entry.record_id = int(record_id)
                entry.timestamp = int(t_us)
                entry.runner_timestamp = int(t_us)

                # Optional system fields: take from first row (best-effort)
                row0 = df.iloc[t_idxs[0]]
                if "request_id" in df.columns and not _is_nan(row0.get("request_id")):
                    entry.request_id = str(row0.get("request_id"))
                if "as_of_timestamp" in df.columns and not _is_nan(row0.get("as_of_timestamp")):
                    as_us = int(_datetime_series_to_us(pd.Series([row0.get("as_of_timestamp")]))[0])  # type: ignore[index]
                    entry.as_of_timestamp = as_us

                for i in t_idxs:
                    row = df.iloc[i]
                    sym_key = str(symbols.iloc[i]).strip()
                    if not sym_key.isdigit():
                        raise ConfigurationError(f"symbol must be numeric symbol_id after mapping, got {sym_key!r}")
                    sym_id = int(sym_key)

                    # Build features list for this row
                    features: List[float]
                    if use_features_col:
                        fv = row.get("features")
                        if isinstance(fv, (list, tuple)):
                            features = [float(x) if x is not None else float("nan") for x in fv]
                        elif isinstance(fv, str):
                            features = _parse_feature_list_str(fv)
                        elif _is_nan(fv):
                            features = []
                        else:
                            raise ValueError("features column must be list/tuple or string like '[1.0, nan, ...]'")
                    else:
                        features = []
                        for name in feature_names:
                            if name in df.columns:
                                v = row.get(name)
                                if _is_nan(v):
                                    features.append(float("nan"))
                                else:
                                    try:
                                        features.append(float(v))
                                    except Exception:
                                        features.append(float("nan"))
                            else:
                                features.append(float("nan"))

                    sv = entry.symbol_values.add()
                    sv.symbol_id = int(sym_id)
                    if encrypt:
                        sv.encrypted_payload = _encrypt_features_aes256gcm_record_key(str(encryption_key), features)
                    else:
                        sv.values.items.extend(features)

                b = entry.SerializeToString()
                buf += _length_prefix_be(len(b))
                buf += b

            url = f"{api_base}{HISTORY_UPLOAD_PATH}"
            # Important: backend expects symbol as SYMBOL NAME for filtering; omit it to upload all symbols in the stream.
            params = {"record_id": int(record_id)}

            # Update overall progress (per-minute request)
            pct = ((done_groups) / max(total_groups, 1)) * 100.0
            progress.update(
                f"Uploading {_format_progress_bar(pct=pct)} {done_groups}/{total_groups} "
                f"sent={_format_bytes(sent_bytes)}"
            )

            resp = _request_with_retries(
                "POST",
                url,
                headers=headers,
                params=params,
                data=bytes(buf),
                timeout=int(DEFAULT_TIMEOUT_S),
                retries=5,
            )
            if resp.status_code != 200:
                raise RequestError(f"Historical upload failed: HTTP {resp.status_code} {resp.text.strip()[:500]}")

            sent_bytes += len(buf)
            done_groups += 1

        # Final 100% render
        progress.update(
            f"Uploading {_format_progress_bar(pct=100.0)} {done_groups}/{total_groups} sent={_format_bytes(sent_bytes)}"
        )
    finally:
        progress.finish()

