from __future__ import annotations

import hashlib
import math
import os
import struct
import time
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)


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 upload_data(
    record_id: int,
    df: "pd.DataFrame",
    api_key: Optional[str] = None,
    private_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.
        private_key: Optional. Record private 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 (private_key or "").strip():
        raise ConfigurationError("Record is encrypted but private_key was not provided.")

    # Prepare timestamp microseconds + grouping key
    ts_us = _datetime_series_to_us(df["timestamp"])
    symbols = df["symbol"].astype(str)
    minute_bucket = (ts_us // 60_000_000).astype("int64")

    # Build per-(symbol, minute) indices
    groups: Dict[_MinuteKey, List[int]] = {}
    for i, (sym, mb) in enumerate(zip(symbols.tolist(), minute_bucket.tolist())):
        key = _MinuteKey(symbol=str(sym), minute_utc=int(mb))
        groups.setdefault(key, []).append(i)

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

    # Process in sorted order for determinism
    for key in sorted(groups.keys(), key=lambda k: (k.symbol, k.minute_utc)):
        idxs = groups[key]

        # Enforce "each request = 1 minute of data"
        # (key already buckets by minute)
        sym = key.symbol

        buf = bytearray()
        for i in idxs:
            row = df.iloc[i]
            entry = data_entry_pb2.DataEntry()
            entry.record_id = int(record_id)
            entry.symbol = str(sym)

            t_us = int(ts_us.iloc[i])
            entry.timestamp = t_us
            # Keep runner_timestamp aligned with timestamp for historical uploads
            entry.runner_timestamp = t_us

            # Optional system fields
            if "request_id" in df.columns and not _is_nan(row.get("request_id")):
                entry.request_id = str(row.get("request_id"))
            if "batch_id" in df.columns and not _is_nan(row.get("batch_id")):
                entry.batch_id = int(row.get("batch_id"))
            if "as_of_timestamp" in df.columns and not _is_nan(row.get("as_of_timestamp")):
                as_us = int(_datetime_series_to_us(pd.Series([row.get("as_of_timestamp")]))[0])  # type: ignore[index]
                entry.as_of_timestamp = as_us

            # Build features list
            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"))

            if encrypt:
                entry.encrypted_payload = _encrypt_features_aes256gcm_record_key(str(private_key), features)
            else:
                entry.features.extend(features)

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

        url = f"{api_base}{HISTORY_UPLOAD_PATH}"
        params = {"record_id": int(record_id), "symbol": sym}
        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]}")

