"""
TradeStation OAuth2 Authentication Module.

Implements Auth Code Flow with local HTTP callback server.
Handles token exchange, refresh, and persistent storage.
"""

import json
import logging
import secrets
import time
import webbrowser
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from urllib.parse import parse_qs, urlencode, urlparse

import httpx

logger = logging.getLogger(__name__)

# TradeStation OAuth2 endpoints
AUTH_URL = "https://signin.tradestation.com/authorize"
TOKEN_URL = "https://signin.tradestation.com/oauth/token"
AUDIENCE = "https://api.tradestation.com"

# Default scopes
DEFAULT_SCOPES = "openid offline_access profile MarketData ReadAccount Trade Matrix OptionSpreads"

# Token storage location
TOKEN_FILE = Path.home() / ".tradestation_mcp_tokens.json"

# Access tokens expire in 1200 seconds (20 minutes).
# Refresh 2 minutes early to avoid mid-request expiry.
TOKEN_REFRESH_BUFFER = 120


class _OAuthCallbackHandler(BaseHTTPRequestHandler):
    """HTTP handler that captures the OAuth authorization code from the redirect."""

    auth_code: str | None = None
    auth_error: str | None = None
    expected_state: str | None = None

    def do_GET(self):
        parsed = urlparse(self.path)
        params = parse_qs(parsed.query)

        if "error" in params:
            _OAuthCallbackHandler.auth_error = params["error"][0]
            error_desc = params.get("error_description", ["Authorization failed"])[0]
            self.send_response(400)
            self.send_header("Content-Type", "text/html")
            self.end_headers()
            self.wfile.write(
                f"<html><body><h2>Authorization Failed</h2><p>{error_desc}</p>"
                f"<p>You can close this window.</p></body></html>".encode()
            )
        elif "code" in params:
            received_state = params.get("state", [None])[0]
            if _OAuthCallbackHandler.expected_state and received_state != _OAuthCallbackHandler.expected_state:
                _OAuthCallbackHandler.auth_error = "state_mismatch"
                self.send_response(400)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(
                    b"<html><body><h2>State Mismatch</h2>"
                    b"<p>CSRF state validation failed. Please try again.</p></body></html>"
                )
            else:
                _OAuthCallbackHandler.auth_code = params["code"][0]
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(
                    b"<html><body><h2>Authorization Successful!</h2>"
                    b"<p>You can close this window and return to your application.</p></body></html>"
                )
        else:
            self.send_response(400)
            self.send_header("Content-Type", "text/html")
            self.end_headers()
            self.wfile.write(b"<html><body><h2>Invalid request</h2></body></html>")

    def log_message(self, format, *args):
        # Silence HTTP server logs — they go to stdout and break STDIO transport
        logger.debug(format, *args)


class TokenManager:
    """Manages OAuth2 tokens for the TradeStation API.

    Handles the full auth code flow, token exchange, automatic refresh,
    and persistent storage of tokens to disk.
    """

    def __init__(
        self,
        client_id: str,
        client_secret: str,
        redirect_port: int = 3000,
    ):
        self.client_id = client_id
        self.client_secret = client_secret
        self.redirect_port = redirect_port
        self.redirect_uri = f"http://localhost:{redirect_port}"

        self._access_token: str | None = None
        self._refresh_token: str | None = None
        self._expires_at: float = 0.0
        self._scopes: str = ""

        # Try to load saved tokens on init
        self._load_tokens()

    # ── Token Persistence ──────────────────────────────────────────

    def _load_tokens(self):
        """Load tokens from the persistent file."""
        if TOKEN_FILE.exists():
            try:
                data = json.loads(TOKEN_FILE.read_text(encoding="utf-8"))
                self._access_token = data.get("access_token")
                self._refresh_token = data.get("refresh_token")
                self._expires_at = data.get("expires_at", 0.0)
                self._scopes = data.get("scopes", "")
                logger.info("Loaded saved tokens from %s", TOKEN_FILE)
            except (json.JSONDecodeError, OSError) as e:
                logger.warning("Failed to load saved tokens: %s", e)

    def _save_tokens(self):
        """Persist current tokens to disk."""
        data = {
            "access_token": self._access_token,
            "refresh_token": self._refresh_token,
            "expires_at": self._expires_at,
            "scopes": self._scopes,
        }
        try:
            TOKEN_FILE.write_text(json.dumps(data, indent=2), encoding="utf-8")
            logger.info("Saved tokens to %s", TOKEN_FILE)
        except OSError as e:
            logger.warning("Failed to save tokens: %s", e)

    # ── OAuth Browser Flow ─────────────────────────────────────────

    def _run_auth_code_flow(self) -> dict:
        """Execute the full OAuth2 Authorization Code Flow.

        1. Start a local HTTP server to receive the callback
        2. Open the user's browser to TradeStation's login page
        3. Wait for the auth code redirect
        4. Exchange the code for tokens

        Returns:
            Token response dict with access_token, refresh_token, etc.

        Raises:
            RuntimeError: If authentication fails at any step.
        """
        state = secrets.token_urlsafe(32)
        _OAuthCallbackHandler.auth_code = None
        _OAuthCallbackHandler.auth_error = None
        _OAuthCallbackHandler.expected_state = state

        # Build authorization URL
        params = {
            "response_type": "code",
            "client_id": self.client_id,
            "redirect_uri": self.redirect_uri,
            "audience": AUDIENCE,
            "scope": DEFAULT_SCOPES,
            "state": state,
            "prompt": "login",
        }
        auth_url = f"{AUTH_URL}?{urlencode(params)}"

        # Start local callback server
        server = HTTPServer(("127.0.0.1", self.redirect_port), _OAuthCallbackHandler)
        server.timeout = 300  # 5 minute timeout waiting for callback

        logger.info("Starting OAuth callback server on port %d", self.redirect_port)
        logger.info("Opening browser for TradeStation login...")

        # Open browser in a thread so we don't block
        webbrowser.open(auth_url)

        # Wait for the callback (blocking)
        while _OAuthCallbackHandler.auth_code is None and _OAuthCallbackHandler.auth_error is None:
            server.handle_request()

        server.server_close()

        if _OAuthCallbackHandler.auth_error:
            raise RuntimeError(f"OAuth authorization failed: {_OAuthCallbackHandler.auth_error}")

        auth_code = _OAuthCallbackHandler.auth_code
        if not auth_code:
            raise RuntimeError("No authorization code received")

        logger.info("Authorization code received, exchanging for tokens...")

        # Exchange auth code for tokens
        return self._exchange_code(auth_code)

    def _exchange_code(self, code: str) -> dict:
        """Exchange an authorization code for access and refresh tokens."""
        payload = {
            "grant_type": "authorization_code",
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "code": code,
            "redirect_uri": self.redirect_uri,
        }

        with httpx.Client() as client:
            resp = client.post(
                TOKEN_URL,
                data=payload,
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                timeout=30.0,
            )

        if resp.status_code != 200:
            raise RuntimeError(f"Token exchange failed ({resp.status_code}): {resp.text}")

        return resp.json()

    # ── Token Refresh ──────────────────────────────────────────────

    def _refresh_access_token(self) -> dict:
        """Use the refresh token to obtain a new access token."""
        if not self._refresh_token:
            raise RuntimeError("No refresh token available. Re-authentication required.")

        payload = {
            "grant_type": "refresh_token",
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "refresh_token": self._refresh_token,
        }

        with httpx.Client() as client:
            resp = client.post(
                TOKEN_URL,
                data=payload,
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                timeout=30.0,
            )

        if resp.status_code != 200:
            raise RuntimeError(f"Token refresh failed ({resp.status_code}): {resp.text}")

        return resp.json()

    # ── Token Application ──────────────────────────────────────────

    def _apply_token_response(self, token_data: dict):
        """Store tokens from an OAuth token response."""
        self._access_token = token_data["access_token"]
        self._scopes = token_data.get("scope", "")
        expires_in = token_data.get("expires_in", 1200)
        self._expires_at = time.time() + expires_in

        # Refresh token may or may not be in the response
        if "refresh_token" in token_data:
            self._refresh_token = token_data["refresh_token"]

        self._save_tokens()
        logger.info(
            "Tokens updated. Access token expires in %d seconds.",
            expires_in,
        )

    # ── Public Interface ───────────────────────────────────────────

    def is_token_valid(self) -> bool:
        """Check if the current access token is still valid (with buffer)."""
        return self._access_token is not None and time.time() < (self._expires_at - TOKEN_REFRESH_BUFFER)

    def has_refresh_token(self) -> bool:
        """Check if a refresh token is available."""
        return self._refresh_token is not None

    def ensure_authenticated(self):
        """Ensure we have a valid access token, refreshing or re-authenticating as needed.

        Call sequence:
        1. If access token is valid → return immediately
        2. If refresh token exists → try refreshing
        3. If refresh fails or no refresh token → run full browser auth flow
        """
        if self.is_token_valid():
            logger.debug("Access token is still valid.")
            return

        if self.has_refresh_token():
            try:
                logger.info("Access token expired, refreshing...")
                token_data = self._refresh_access_token()
                self._apply_token_response(token_data)
                return
            except RuntimeError as e:
                logger.warning("Token refresh failed: %s. Will re-authenticate.", e)

        # Full auth flow needed
        logger.info("Starting full OAuth authentication flow...")
        token_data = self._run_auth_code_flow()
        self._apply_token_response(token_data)

    def get_access_token(self) -> str:
        """Get a valid access token, refreshing if necessary.

        Returns:
            A valid bearer token string.
        """
        self.ensure_authenticated()
        return self._access_token  # type: ignore[return-value]
