"""
auth.py — User authentication for llm-guard-kit SaaS.

Provides:
  - Email/password registration and login
  - Google OAuth2 flow (no SDK — plain HTTPS calls via httpx)
  - JWT session tokens (HS256, 72-hour expiry)
  - Plan enforcement helpers (Free: 500 chains/mo, Pro: unlimited + judge)

Environment variables:
  JWT_SECRET          — MUST be set in production (random if not set)
  GOOGLE_CLIENT_ID    — Google OAuth app client ID
  GOOGLE_CLIENT_SECRET — Google OAuth app client secret

Usage:
  from qppg_service.auth import make_auth_router, get_current_user, check_plan_limits
  app.include_router(make_auth_router(store), prefix="/auth", tags=["auth"])
"""

from __future__ import annotations

import hashlib
import hmac
import os
import secrets
from datetime import datetime, timedelta, timezone

import jwt
from fastapi import APIRouter, Depends, Header, HTTPException
from pydantic import BaseModel

from qppg_service.store import ChainStore

# ── Constants ──────────────────────────────────────────────────────────────────

JWT_SECRET: str = os.environ.get("JWT_SECRET", secrets.token_hex(32))
JWT_ALGORITHM = "HS256"
JWT_EXPIRY_HOURS = 72

FREE_MONTHLY_LIMIT = 500

# ── Password hashing (stdlib PBKDF2 — no extra deps) ──────────────────────────

_HASH_ITERATIONS = 600_000  # NIST 2023 recommendation for PBKDF2-HMAC-SHA256


def hash_password(password: str) -> str:
    """Return a salted PBKDF2-SHA256 hash: 'pbkdf2:sha256:<iters>$<salt>$<hash>'"""
    salt = secrets.token_hex(16)
    dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), _HASH_ITERATIONS)
    return f"pbkdf2:sha256:{_HASH_ITERATIONS}${salt}${dk.hex()}"


def verify_password(password: str, stored_hash: str) -> bool:
    """Verify a password against a stored PBKDF2 hash.

    Expected format: pbkdf2:sha256:<iterations>$<salt>$<hex_dk>
    """
    try:
        _, _algo, rest = stored_hash.split(":", 2)
        iters_str, salt, dk_hex = rest.split("$", 2)
        iters = int(iters_str)
        dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), iters)
        return hmac.compare_digest(dk.hex(), dk_hex)
    except Exception:
        return False

GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v3/userinfo"

# ── JWT helpers ────────────────────────────────────────────────────────────────


def make_token(user_id: int) -> str:
    """Create a signed JWT for the given user ID."""
    payload = {
        "sub": str(user_id),
        "exp": datetime.now(tz=timezone.utc) + timedelta(hours=JWT_EXPIRY_HOURS),
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)


def decode_token(token: str) -> int | None:
    """Decode a JWT; returns user_id int or None on any failure."""
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        return int(payload["sub"])
    except Exception:
        return None


# ── FastAPI dependency ─────────────────────────────────────────────────────────


def _get_current_user_factory(store: ChainStore):
    """Returns a FastAPI dependency that validates Bearer JWT and returns the user dict."""

    def get_current_user(authorization: str | None = Header(default=None)) -> dict:
        if not authorization or not authorization.startswith("Bearer "):
            raise HTTPException(401, "Missing or invalid Authorization header")
        token = authorization.removeprefix("Bearer ").strip()
        user_id = decode_token(token)
        if user_id is None:
            raise HTTPException(401, "Invalid or expired token")
        user = store.get_user_by_id(user_id)
        if not user or not user.get("is_active"):
            raise HTTPException(401, "User not found or deactivated")
        return user

    return get_current_user


# ── Plan enforcement ───────────────────────────────────────────────────────────


def check_plan_limits(user: dict, store: ChainStore, want_judge: bool = False) -> None:
    """
    Raises HTTPException if the user has exceeded their plan limits.

    Free plan:
      - No judge scoring (403)
      - 500 chains/month cap (429)
    Pro plan:
      - Unlimited chains, all signals allowed
    """
    if user["plan"] == "free":
        if want_judge:
            raise HTTPException(
                403,
                detail="Judge scoring requires Pro. Upgrade at /pricing.",
            )
        usage = store.get_monthly_usage(user["id"])
        if usage >= FREE_MONTHLY_LIMIT:
            raise HTTPException(
                429,
                detail=f"Free plan limit ({FREE_MONTHLY_LIMIT} chains/mo). Upgrade at /pricing.",
            )


# ── Pydantic models ────────────────────────────────────────────────────────────


class RegisterRequest(BaseModel):
    email: str
    password: str
    name: str = ""


class LoginRequest(BaseModel):
    email: str
    password: str


class GoogleCallbackRequest(BaseModel):
    code: str
    redirect_uri: str


# ── Router factory ─────────────────────────────────────────────────────────────


def make_auth_router(store: ChainStore) -> APIRouter:
    """Build and return the auth APIRouter bound to the given ChainStore."""

    router = APIRouter()
    get_current_user = _get_current_user_factory(store)

    # ── Register ──────────────────────────────────────────────────────────────

    @router.post("/register")
    def register(body: RegisterRequest):
        """Create a new account with email + password."""
        if store.get_user_by_email(body.email):
            raise HTTPException(409, "An account with this email already exists")

        password_hash = hash_password(body.password)
        user = store.create_user(
            email=body.email,
            password_hash=password_hash,
            name=body.name,
        )
        return {
            "access_token": make_token(user["id"]),
            "user": _public_user(user),
        }

    # ── Login ─────────────────────────────────────────────────────────────────

    @router.post("/login")
    def login(body: LoginRequest):
        """Authenticate with email + password."""
        user = store.get_user_by_email(body.email)
        if not user or not user.get("password_hash"):
            raise HTTPException(401, "Invalid email or password")
        if not verify_password(body.password, user["password_hash"]):
            raise HTTPException(401, "Invalid email or password")
        if not user.get("is_active"):
            raise HTTPException(401, "Account is deactivated")

        return {
            "access_token": make_token(user["id"]),
            "user": _public_user(user),
        }

    # ── Google OAuth — initiate ───────────────────────────────────────────────

    @router.get("/google")
    def google_auth_url(redirect_uri: str):
        """Return the Google OAuth authorization URL for the frontend to redirect to."""
        client_id = os.environ.get("GOOGLE_CLIENT_ID", "")
        if not client_id:
            raise HTTPException(500, "GOOGLE_CLIENT_ID not configured")

        state = secrets.token_urlsafe(16)
        params = {
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "response_type": "code",
            "scope": "openid email profile",
            "state": state,
            "access_type": "offline",
            "prompt": "select_account",
        }
        from urllib.parse import urlencode
        url = GOOGLE_AUTH_URL + "?" + urlencode(params)
        return {"url": url}

    # ── Google OAuth — callback ───────────────────────────────────────────────

    @router.post("/google/callback")
    async def google_callback(body: GoogleCallbackRequest):
        """Exchange Google auth code for an access token; create/login user."""
        try:
            import httpx
        except ImportError:
            raise HTTPException(500, "httpx not installed — run: pip install 'llm-guard-kit[saas]'")

        client_id = os.environ.get("GOOGLE_CLIENT_ID", "")
        client_secret = os.environ.get("GOOGLE_CLIENT_SECRET", "")
        if not client_id or not client_secret:
            raise HTTPException(500, "Google OAuth not configured")

        # Exchange authorization code for tokens
        async with httpx.AsyncClient() as client:
            token_resp = await client.post(
                GOOGLE_TOKEN_URL,
                data={
                    "code": body.code,
                    "client_id": client_id,
                    "client_secret": client_secret,
                    "redirect_uri": body.redirect_uri,
                    "grant_type": "authorization_code",
                },
            )

        if token_resp.status_code != 200:
            raise HTTPException(400, "Failed to exchange Google auth code")

        token_data = token_resp.json()
        id_token_str = token_data.get("id_token", "")
        if not id_token_str:
            raise HTTPException(400, "Google did not return an id_token")

        # Decode id_token without signature verification (we trust Google's HTTPS)
        import base64
        import json as _json

        parts = id_token_str.split(".")
        if len(parts) < 2:
            raise HTTPException(400, "Malformed Google id_token")
        # Pad to valid base64
        payload_b64 = parts[1] + "=" * (4 - len(parts[1]) % 4)
        claims = _json.loads(base64.urlsafe_b64decode(payload_b64))

        google_id = claims.get("sub")
        email = claims.get("email", "")
        name = claims.get("name", "")

        if not google_id or not email:
            raise HTTPException(400, "Google token missing sub or email")

        # Find or create user
        user = store.get_user_by_google_id(google_id)
        if not user:
            # Check if email already registered with password
            existing = store.get_user_by_email(email)
            if existing:
                # Link Google ID to existing account
                with store._conn() as conn:
                    conn.execute(
                        "UPDATE users SET google_id=? WHERE id=?",
                        (google_id, existing["id"]),
                    )
                user = store.get_user_by_id(existing["id"])
            else:
                user = store.create_user(email=email, google_id=google_id, name=name)

        return {
            "access_token": make_token(user["id"]),
            "user": _public_user(user),
        }

    # ── Me ────────────────────────────────────────────────────────────────────

    @router.get("/me")
    def me(user: dict = Depends(get_current_user)):
        """Return the current user's profile."""
        return {
            **_public_user(user),
            "chains_this_month": store.get_monthly_usage(user["id"]),
            "created_at": user["created_at"],
        }

    # ── Logout ────────────────────────────────────────────────────────────────

    @router.post("/logout")
    def logout():
        """Stateless logout — frontend must discard the token."""
        return {"ok": True}

    return router


# ── Internal helpers ───────────────────────────────────────────────────────────


def _public_user(user: dict) -> dict:
    """Return only the fields safe to expose in API responses."""
    return {
        "id": user["id"],
        "email": user["email"],
        "name": user.get("name", ""),
        "plan": user.get("plan", "free"),
    }
