from __future__ import annotations

import base64
import datetime as dt
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional
from urllib.parse import urljoin

import requests
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding


@dataclass(frozen=True)
class KalshiConfig:
    env: str  # "demo" or "prod"


class KalshiClient:
    """
    Minimal Kalshi REST client (public market data + authenticated trading).
    Auth scheme: KALSHI-ACCESS-KEY, KALSHI-ACCESS-TIMESTAMP (ms), KALSHI-ACCESS-SIGNATURE
    where signature is RSA-PSS-SHA256 of: timestamp + METHOD + path_without_query. :contentReference[oaicite:1]{index=1}
    """

    def __init__(self, cfg: KalshiConfig, session: Optional[requests.Session] = None):
        self.cfg = cfg
        self.session = session or requests.Session()

        # Base URL differs for demo vs prod per Kalshi quick-start docs. :contentReference[oaicite:2]{index=2}
        if cfg.env.lower() == "demo":
            self.base_url = "https://demo-api.kalshi.co"
        elif cfg.env.lower() in ("prod", "production"):
            self.base_url = "https://api.kalshi.com"
        else:
            raise ValueError("KALSHI_ENV must be 'demo' or 'prod'")

        # Auth fields are optional in paper mode; only required when placing orders
        self.api_key_id = os.getenv("KALSHI_API_KEY_ID")
        self.private_key_path = os.getenv("KALSHI_PRIVATE_KEY_PATH")
        self._private_key = None

    def _load_private_key(self) -> None:
        if self._private_key is not None:
            return
        if not self.private_key_path:
            raise RuntimeError("Missing KALSHI_PRIVATE_KEY_PATH (did you source .env?)")
        with open(self.private_key_path, "rb") as f:
            self._private_key = serialization.load_pem_private_key(f.read(), password=None)

    def _timestamp_ms(self) -> str:
        return str(int(dt.datetime.now(dt.timezone.utc).timestamp() * 1000))

    def _sign(self, timestamp_ms: str, method: str, path: str) -> str:
        """
        Sign timestamp + METHOD + path_without_query using RSA-PSS(SHA256), return base64 signature.
        :contentReference[oaicite:3]{index=3}
        """
        self._load_private_key()
        assert self._private_key is not None

        path_wo_query = path.split("?", 1)[0]
        msg = f"{timestamp_ms}{method.upper()}{path_wo_query}".encode("utf-8")

        sig = self._private_key.sign(
            msg,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.DIGEST_LENGTH,
            ),
            hashes.SHA256(),
        )
        return base64.b64encode(sig).decode("utf-8")

    def _auth_headers(self, method: str, path: str) -> Dict[str, str]:
        if not self.api_key_id:
            raise RuntimeError("Missing KALSHI_API_KEY_ID (did you source .env?)")
        ts = self._timestamp_ms()
        sig = self._sign(ts, method, path)
        return {
            "KALSHI-ACCESS-KEY": self.api_key_id,
            "KALSHI-ACCESS-TIMESTAMP": ts,
            "KALSHI-ACCESS-SIGNATURE": sig,
        }

    def get_public(self, path: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        url = urljoin(self.base_url, path)
        r = self.session.get(url, params=params, timeout=20)
        r.raise_for_status()
        return r.json()

    def get_authed(self, path: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        url = urljoin(self.base_url, path)
        headers = self._auth_headers("GET", path)
        r = self.session.get(url, headers=headers, params=params, timeout=20)
        r.raise_for_status()
        return r.json()

    def post_authed(self, path: str, json_body: Dict[str, Any]) -> Dict[str, Any]:
        url = urljoin(self.base_url, path)
        headers = self._auth_headers("POST", path)
        headers["Content-Type"] = "application/json"
        r = self.session.post(url, headers=headers, json=json_body, timeout=20)
        r.raise_for_status()
        return r.json()

    # Convenience wrappers
    def list_open_markets(self, limit: int = 1000) -> Dict[str, Any]:
        # public market data endpoint per quick-start :contentReference[oaicite:4]{index=4}
        return self.get_public("/trade-api/v2/markets", params={"status": "open", "limit": str(limit)})

    def get_market(self, ticker: str) -> Dict[str, Any]:
        return self.get_public(f"/trade-api/v2/markets/{ticker}")

    def create_order(self, order_data: Dict[str, Any]) -> Dict[str, Any]:
        # POST /trade-api/v2/portfolio/orders per quick-start :contentReference[oaicite:5]{index=5}
        return self.post_authed("/trade-api/v2/portfolio/orders", order_data)