"""JWT token operations only."""
from __future__ import annotations
from datetime import datetime, timedelta, timezone
from typing import Optional
from uuid import UUID

from jose import JWTError, jwt
from pydantic import BaseModel, ValidationError

from src.config import get_settings

settings = get_settings()


class TokenPayload(BaseModel):
    """JWT payload structure."""

    sub: str  # user_id as string
    exp: datetime
    iat: datetime
    type: str  # "access" | "refresh" | "mfa_pending" | "desktop_session"
    session_id: Optional[str] = None


def create_access_token(user_id: UUID, session_id: Optional[UUID] = None) -> str:
    """Create JWT access token with short TTL."""
    now = datetime.now(timezone.utc)
    expire = now + timedelta(minutes=settings.jwt_access_ttl_minutes)
    payload = {
        "sub": str(user_id),
        "exp": expire,
        "iat": now,
        "type": "access",
        "session_id": str(session_id) if session_id else None,
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


def create_refresh_token(user_id: UUID, session_id: UUID) -> str:
    """Create JWT refresh token with long TTL."""
    now = datetime.now(timezone.utc)
    expire = now + timedelta(days=settings.jwt_refresh_ttl_days)
    payload = {
        "sub": str(user_id),
        "exp": expire,
        "iat": now,
        "type": "refresh",
        "session_id": str(session_id),
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


def decode_token(token: str) -> Optional[TokenPayload]:
    """Decode and validate JWT token. Returns None if invalid."""
    try:
        payload = jwt.decode(token, settings.jwt_secret, algorithms=["HS256"])
        return TokenPayload(
            sub=payload["sub"],
            exp=datetime.fromtimestamp(payload["exp"], tz=timezone.utc),
            iat=datetime.fromtimestamp(payload["iat"], tz=timezone.utc),
            type=payload["type"],
            session_id=payload.get("session_id"),
        )
    except JWTError:
        return None


def get_user_id_from_token(token: str) -> Optional[UUID]:
    """Extract user ID from access token. Returns None if invalid."""
    payload = decode_token(token)
    if payload and payload.type == "access":
        try:
            return UUID(payload.sub)
        except ValueError:
            return None
    return None


def create_mfa_pending_token(user_id: UUID, session_id: UUID) -> str:
    """Create short-lived token for MFA completion (5 min)."""
    now = datetime.now(timezone.utc)
    expire = now + timedelta(minutes=5)
    payload = {
        "sub": str(user_id),
        "exp": expire,
        "iat": now,
        "type": "mfa_pending",
        "session_id": str(session_id),
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


def decode_mfa_pending_token(token: str) -> Optional[tuple[UUID, UUID]]:
    """Decode MFA pending token. Returns (user_id, session_id) or None."""
    payload = decode_token(token)
    if payload and payload.type == "mfa_pending" and payload.session_id:
        try:
            return UUID(payload.sub), UUID(payload.session_id)
        except ValueError:
            return None
    return None


class RunGrantPayload(BaseModel):
    """Payload for a short-lived per-run grant issued to a paired desktop device.

    ``sub`` is the device_id (matches TokenPayload's convention of "subject = who this is
    issued to"). The durable ``device_secret`` itself is rarely transmitted; Core mints one of
    these per run instead, so a leaked grant has a small blast radius and short useful life.
    """

    sub: str  # device_id as string
    run_id: str
    project_id: str
    workspace_id: str
    exp: datetime
    iat: datetime
    type: str  # "run_grant"


class CodegenActionGrantPayload(BaseModel):
    """Short-lived grant embedded in a Core-signed managed Actions payload."""

    run_id: str
    project_id: str
    workspace_id: str
    exp: datetime
    iat: datetime
    type: str  # "codegen_action_grant"


# Hardcoded rather than settings-driven: no other TTL in this codebase is exposed via config
# except the user-facing access/refresh tokens; this is an internal device<->Core grant.
RUN_GRANT_TTL_MINUTES = 20

# The GitHub Actions codegen job itself runs with `timeout-minutes: 30` (see
# actions_pipeline.py), and the grant is minted once at dispatch time, before the
# job even starts, then reused for every progress/cloud-cutover/completion callback
# for the rest of that run. It must outlive the job's own timeout (plus queueing,
# retry backoff, and continuation re-dispatch) or completion callbacks start
# failing closed with 403 right as a run finishes -- which is silent data loss,
# not a safe failure. Generous on purpose.
CODEGEN_ACTION_GRANT_TTL_MINUTES = 120


def create_run_grant_token(
    device_id: UUID,
    run_id: UUID,
    project_id: UUID,
    workspace_id: UUID,
) -> str:
    """Create a short-lived run-grant JWT authorizing one device to report progress/completion
    for one specific run (see require_run_grant_or_actions_token)."""
    now = datetime.now(timezone.utc)
    expire = now + timedelta(minutes=RUN_GRANT_TTL_MINUTES)
    payload = {
        "sub": str(device_id),
        "run_id": str(run_id),
        "project_id": str(project_id),
        "workspace_id": str(workspace_id),
        "exp": expire,
        "iat": now,
        "type": "run_grant",
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


DESKTOP_SESSION_TTL_MINUTES = 30


def create_desktop_session_token(user_id: UUID) -> str:
    """Short-lived session for a desktop app after browser-led login exchange.

    Distinct from web ``access`` tokens so the Mac app cannot call cookie-shaped Dash APIs,
    and distinct from ``device_secret`` which is only issued after pairing.
    """
    now = datetime.now(timezone.utc)
    expire = now + timedelta(minutes=DESKTOP_SESSION_TTL_MINUTES)
    payload = {
        "sub": str(user_id),
        "exp": expire,
        "iat": now,
        "type": "desktop_session",
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


def decode_desktop_session_token(token: str) -> Optional[UUID]:
    """Return user_id if this is a valid desktop session JWT."""
    payload = decode_token(token)
    if not payload or payload.type != "desktop_session":
        return None
    try:
        return UUID(payload.sub)
    except ValueError:
        return None


def decode_run_grant_token(token: str) -> Optional[RunGrantPayload]:
    """Decode and validate a run-grant JWT. Returns None if invalid, expired, or wrong type."""
    try:
        payload = jwt.decode(token, settings.jwt_secret, algorithms=["HS256"])
    except JWTError:
        return None
    if payload.get("type") != "run_grant":
        return None
    try:
        return RunGrantPayload(
            sub=payload["sub"],
            run_id=payload["run_id"],
            project_id=payload["project_id"],
            workspace_id=payload["workspace_id"],
            exp=datetime.fromtimestamp(payload["exp"], tz=timezone.utc),
            iat=datetime.fromtimestamp(payload["iat"], tz=timezone.utc),
            type=payload["type"],
        )
    except (KeyError, ValueError, ValidationError):
        return None


def create_codegen_action_grant_token(
    run_id: UUID,
    project_id: UUID,
    workspace_id: UUID,
) -> str:
    """Create a short-lived grant for the P2HS-managed GitHub Actions runner.

    Unlike the legacy repository secret, this grant is scoped to exactly one run,
    so a stale or incorrectly rotated Actions secret cannot block that run's
    Core callbacks.
    """
    now = datetime.now(timezone.utc)
    expire = now + timedelta(minutes=CODEGEN_ACTION_GRANT_TTL_MINUTES)
    payload = {
        "run_id": str(run_id),
        "project_id": str(project_id),
        "workspace_id": str(workspace_id),
        "exp": expire,
        "iat": now,
        "type": "codegen_action_grant",
    }
    return jwt.encode(payload, settings.jwt_secret, algorithm="HS256")


def decode_codegen_action_grant_token(token: str) -> Optional[CodegenActionGrantPayload]:
    """Decode a managed GitHub Actions grant, returning None when invalid."""
    try:
        payload = jwt.decode(token, settings.jwt_secret, algorithms=["HS256"])
        grant = CodegenActionGrantPayload.model_validate(payload)
    except (JWTError, ValidationError):
        return None
    if grant.type != "codegen_action_grant":
        return None
    return grant