from __future__ import annotations

import logging
import math
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Literal, Optional

from fastapi import APIRouter, Body, Depends, HTTPException, Query
from pydantic import BaseModel, Field
from sqlalchemy import and_, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from ...auth_security import AdminContext, require_admin
from ...db import get_session
from ...models import (
    Account,
    AccountEntitlements,
    AccountUsage,
    BetaSubmission,
    TrialEmailSend,
    TrialEmailTemplate,
    UsageQuotaModel,
    User,
)
from ...services.admin_audit import log_admin_action
from ...services.admin_surface_signals import (
    DEFAULT_PRODUCT_SURFACE,
    SUPPORTED_PRODUCT_SURFACES,
    add_keyword_surface_reasons,
    classify_account_identity,
    collect_account_surface_signals,
    empty_surface_reasons,
    infer_product_surface,
    normalize_product_surface,
    SurfaceSignals,
)
from .admin_customers import _apply_tier_update, _resolve_entitlements_from_limits

log = logging.getLogger(__name__)

router = APIRouter(prefix="/admin/trials", tags=["Admin"])

TRIAL_STATUS_FILTERS = {
    "all",
    "pending",
    "invited",
    "provisioning",
    "active",
    "expired",
    "revoked",
    "canceled",
}

TRIAL_STATUS_PRIORITY = {
    "provisioning": 0,
    "pending": 1,
    "invited": 2,
    "active": 3,
    "expired": 4,
    "revoked": 5,
    "canceled": 6,
}


def _primary_user_subquery(column_name: str):
    return (
        select(getattr(User, column_name))
        .where(User.account_id == Account.id)
        .order_by(User.created_at.asc(), User.id.asc())
        .limit(1)
        .scalar_subquery()
    )


def _user_count_subquery():
    return (
        select(func.count(User.id))
        .where(User.account_id == Account.id)
        .scalar_subquery()
    )


class TrialSurfacePayload(BaseModel):
    product_surface: Optional[Literal["gremlin", "sheets", "deal_desk"]] = None


class TrialExtendPayload(TrialSurfacePayload):
    extend_days: int = Field(default=30, ge=1, le=365)


class TrialRevokePayload(TrialSurfacePayload):
    reason: Optional[str] = None


class TrialUpgradePayload(TrialSurfacePayload):
    tier: Literal["unleashed"] = "unleashed"


def _iso_or_none(value: Optional[datetime]) -> Optional[str]:
    normalized = _ensure_utc(value)
    return normalized.isoformat() if normalized else None


def _ensure_utc(value: Optional[datetime]) -> Optional[datetime]:
    if value is None:
        return None
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


def _resolved_trial_status(
    *,
    quota_status: Optional[str],
    entitlement_tier: Optional[str],
    expires_at: Optional[datetime],
) -> str:
    normalized_status = (quota_status or "").strip().lower()
    normalized_tier = (entitlement_tier or "free").strip().lower()
    now = datetime.now(timezone.utc)
    normalized_expires_at = _ensure_utc(expires_at)

    if normalized_status == "revoked":
        return "revoked"
    if normalized_status in {"canceled", "cancelled"}:
        return "canceled"
    if normalized_tier == "free" and normalized_expires_at and normalized_expires_at <= now:
        return "expired"
    return "active"


def _days_remaining(expires_at: Optional[datetime]) -> Optional[int]:
    normalized_expires_at = _ensure_utc(expires_at)
    if not normalized_expires_at:
        return None
    delta = normalized_expires_at - datetime.now(timezone.utc)
    return max(0, math.ceil(delta.total_seconds() / 86400))


def _email_domain(value: Optional[str]) -> Optional[str]:
    if not value or "@" not in value:
        return None
    return value.split("@", 1)[1].strip().lower() or None


def _looks_like_email(value: Optional[str]) -> bool:
    return bool(value and "@" in value)


def _lead_identity(*, label: str, reason: str) -> Dict[str, object]:
    return {
        "kind": "lead_only",
        "label": label,
        "reason": reason,
        "is_user_fallback": False,
        "clerk_org_id": None,
        "primary_user_id": None,
        "primary_clerk_user_id": None,
        "user_count": 0,
    }


def _surface_flags_for(surface: str) -> Dict[str, bool]:
    normalized = normalize_product_surface(surface)
    return {
        candidate: candidate == normalized
        for candidate in SUPPORTED_PRODUCT_SURFACES
    }


def _surface_reasons_for(surface: str, *reasons: str) -> Dict[str, list[str]]:
    payload = empty_surface_reasons()
    normalized = normalize_product_surface(surface)
    for reason in reasons:
        if reason:
            payload[normalized].append(reason)
    return payload


def _orphan_account_candidate_clause(user_count_subquery):
    return and_(
        Account.id.like("user_%"),
        Account.clerk_org_id.is_(None),
        user_count_subquery == 0,
        Account.name.ilike("%@%"),
        UsageQuotaModel.account_id.is_(None),
        AccountEntitlements.account_id.is_(None),
    )


def _is_orphan_account_row(
    *,
    account_id: str,
    account_name: Optional[str],
    clerk_org_id: Optional[str],
    user_count: int,
    quota_status: Optional[str],
    entitlement_tier: Optional[str],
    trial_started_at: Optional[datetime],
    expires_at: Optional[datetime],
) -> bool:
    return (
        str(account_id).startswith("user_")
        and not clerk_org_id
        and int(user_count or 0) == 0
        and _looks_like_email(account_name)
        and not quota_status
        and not entitlement_tier
        and trial_started_at is None
        and expires_at is None
    )


def _effective_primary_email(
    *,
    primary_email: Optional[str],
    account_name: Optional[str],
) -> Optional[str]:
    if primary_email:
        return primary_email.strip().lower()
    if _looks_like_email(account_name):
        return str(account_name).strip().lower()
    return None


def _search_matches(item: Dict[str, Any], search: Optional[str]) -> bool:
    normalized = (search or "").strip().lower()
    if not normalized:
        return True

    candidates = [
        item.get("id"),
        item.get("account_id"),
        item.get("account_name"),
        item.get("user_id"),
        item.get("primary_email"),
        item.get("email_domain"),
        item.get("sync_source"),
        item.get("status"),
        item.get("record_kind"),
        item.get("status_detail"),
        item.get("identity", {}).get("clerk_org_id"),
    ]
    candidates.extend(item.get("active_surfaces") or [])
    for reasons in (item.get("surface_reasons") or {}).values():
        candidates.extend(reasons)

    return any(
        normalized in str(candidate).lower()
        for candidate in candidates
        if candidate
    )


def _sort_timestamp(value: Optional[str]) -> float:
    if not value:
        return 0.0
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
    except ValueError:
        return 0.0


def _inventory_sort_key(item: Dict[str, Any]) -> tuple[float, float, str]:
    priority = float(TRIAL_STATUS_PRIORITY.get(item.get("status"), 99))
    if item.get("status") == "active":
        days_remaining = item.get("days_remaining")
        return (
            priority,
            float(days_remaining if days_remaining is not None else 10**9),
            -_sort_timestamp(item.get("updated_at")),
        )

    timeline_ts = (
        item.get("invited_at")
        or item.get("submitted_at")
        or item.get("updated_at")
        or item.get("trial_started_at")
    )
    return (
        priority,
        -_sort_timestamp(timeline_ts),
        0.0,
    )


async def _upsert_trial_entitlements(
    db: AsyncSession,
    *,
    account_id: str,
    tier: str,
    source: str,
    monthly_rows: int,
    trial_started_at: Optional[datetime],
    expires_at: Optional[datetime],
) -> AccountEntitlements:
    entitlements = await db.get(AccountEntitlements, account_id)
    now = datetime.now(timezone.utc)

    if not entitlements:
        entitlements = AccountEntitlements(
            account_id=account_id,
            tier=tier,
            addons=[],
            source=source,
            monthly_row_limit=monthly_rows,
            expires_at=expires_at,
        )

    entitlements.tier = tier
    entitlements.source = source
    entitlements.synced_at = now
    entitlements.monthly_row_limit = monthly_rows
    entitlements.addons = list(entitlements.addons or [])
    if trial_started_at is not None or entitlements.trial_started_at is None:
        entitlements.trial_started_at = trial_started_at
    entitlements.expires_at = expires_at

    db.add(entitlements)
    await db.commit()
    await db.refresh(entitlements)
    return entitlements


def _active_surfaces(surface_flags: Dict[str, bool]) -> list[str]:
    return [
        surface
        for surface in SUPPORTED_PRODUCT_SURFACES
        if bool(surface_flags.get(surface))
    ]


def _resolve_primary_surface(
    *,
    surface_flags: Dict[str, bool],
    sync_source: Optional[str],
) -> str:
    active_surfaces = _active_surfaces(surface_flags)
    inferred = infer_product_surface(source=sync_source)
    if inferred in active_surfaces:
        return inferred
    if len(active_surfaces) == 1:
        return active_surfaces[0]
    if active_surfaces:
        return active_surfaces[0]
    return inferred


def _trial_action_block_reason(
    *,
    identity: Dict[str, object],
    active_surfaces: list[str],
) -> Optional[str]:
    if len(active_surfaces) <= 1:
        return None

    if bool(identity.get("is_user_fallback")):
        return (
            "Blocked: this fallback account is active on multiple product surfaces, "
            "and extend/revoke/upgrade are still account-wide mutations. Move the "
            "user into a dedicated Clerk org/account before changing trial state."
        )

    return (
        "Blocked: this account is active on multiple product surfaces, and the "
        "admin trial mutations are account-wide rather than surface-scoped."
    )


def _build_trial_action_source(*, action: str, product_surface: str) -> str:
    normalized_surface = normalize_product_surface(product_surface)
    return f"admin_trial_{action}_{normalized_surface}"


async def _fetch_trial_inventory(
    db: AsyncSession,
    *,
    search: Optional[str] = None,
) -> list[Dict[str, Any]]:
    primary_user_id = _primary_user_subquery("id").label("primary_user_id")
    primary_email = _primary_user_subquery("email").label("primary_email")
    primary_clerk_user_id = _primary_user_subquery("clerk_user_id").label(
        "primary_clerk_user_id"
    )
    primary_user_created_at = _primary_user_subquery("created_at").label(
        "primary_user_created_at"
    )
    user_count_subquery = _user_count_subquery()
    user_count = user_count_subquery.label("user_count")

    query = (
        select(
            Account.id.label("account_id"),
            Account.name.label("account_name"),
            Account.clerk_org_id,
            primary_user_id,
            primary_email,
            primary_clerk_user_id,
            primary_user_created_at,
            user_count,
            UsageQuotaModel.tier.label("quota_tier"),
            UsageQuotaModel.subscription_status.label("quota_status"),
            UsageQuotaModel.tier_sync_source.label("tier_sync_source"),
            UsageQuotaModel.tier_synced_at.label("tier_synced_at"),
            AccountEntitlements.tier.label("entitlement_tier"),
            AccountEntitlements.source.label("entitlement_source"),
            AccountEntitlements.trial_started_at,
            AccountEntitlements.expires_at,
            AccountEntitlements.gremlin_daily_limit,
            AccountEntitlements.addons,
            AccountUsage.gremlin_messages_used,
            AccountUsage.webhook_rows_used,
        )
        .select_from(Account)
        .outerjoin(UsageQuotaModel, UsageQuotaModel.account_id == Account.id)
        .outerjoin(AccountEntitlements, AccountEntitlements.account_id == Account.id)
        .outerjoin(AccountUsage, AccountUsage.account_id == Account.id)
        .where(
            or_(
                AccountEntitlements.trial_started_at.isnot(None),
                AccountEntitlements.expires_at.isnot(None),
                UsageQuotaModel.subscription_status.in_(
                    ("trialing", "revoked", "canceled", "cancelled")
                ),
                _orphan_account_candidate_clause(user_count_subquery),
            )
        )
        .order_by(AccountEntitlements.expires_at.asc().nulls_last(), Account.name.asc())
    )

    rows = (await db.execute(query)).all()
    account_ids = [row.account_id for row in rows]
    primary_emails = {
        row.account_id: _effective_primary_email(
            primary_email=row.primary_email,
            account_name=row.account_name,
        )
        for row in rows
    }
    surface_signals_map = await collect_account_surface_signals(
        db,
        account_ids=account_ids,
        primary_emails_by_account=primary_emails,
    )

    items: list[Dict[str, Any]] = []
    accounted_email_surface_pairs: set[tuple[str, str]] = set()
    for row in rows:
        primary_email_value = _effective_primary_email(
            primary_email=row.primary_email,
            account_name=row.account_name,
        )
        identity = classify_account_identity(
            account_id=row.account_id,
            account_name=row.account_name,
            clerk_org_id=row.clerk_org_id,
            primary_user_id=row.primary_user_id,
            primary_clerk_user_id=row.primary_clerk_user_id,
            user_count=row.user_count or 0,
        )
        orphan_account = _is_orphan_account_row(
            account_id=row.account_id,
            account_name=row.account_name,
            clerk_org_id=row.clerk_org_id,
            user_count=row.user_count or 0,
            quota_status=row.quota_status,
            entitlement_tier=row.entitlement_tier,
            trial_started_at=row.trial_started_at,
            expires_at=row.expires_at,
        )
        signals = surface_signals_map.get(row.account_id) or SurfaceSignals()
        if int(row.gremlin_daily_limit or 0) > 0:
            signals.add(
                "gremlin",
                f"entitlements.gremlin_daily_limit:{int(row.gremlin_daily_limit)}",
            )
        if int(row.gremlin_messages_used or 0) > 0:
            signals.add(
                "gremlin",
                f"account_usage.gremlin_messages_used:{int(row.gremlin_messages_used)}",
            )
        if int(row.webhook_rows_used or 0) > 0:
            signals.add(
                "sheets",
                f"account_usage.webhook_rows_used:{int(row.webhook_rows_used)}",
            )
        add_keyword_surface_reasons(
            signals,
            reason_prefix="source_hint",
            values=(
                row.tier_sync_source,
                row.entitlement_source,
                " ".join(row.addons or []) if row.addons else None,
            ),
        )
        if orphan_account and not any(signals.flags().values()):
            signals.add(
                DEFAULT_PRODUCT_SURFACE,
                "fallback_inference:no explicit surface signals; defaulted to gremlin for orphan account",
            )
        surface_flags = signals.flags()
        active_surfaces = _active_surfaces(surface_flags)
        primary_surface = _resolve_primary_surface(
            surface_flags=surface_flags,
            sync_source=row.tier_sync_source or row.entitlement_source,
        )
        actions_blocked_reason = _trial_action_block_reason(
            identity=identity,
            active_surfaces=active_surfaces,
        )

        status = (
            "provisioning"
            if orphan_account
            else _resolved_trial_status(
                quota_status=row.quota_status,
                entitlement_tier=row.entitlement_tier or row.quota_tier,
                expires_at=row.expires_at,
            )
        )
        resolved_tier = (row.entitlement_tier or row.quota_tier or "free").strip().lower()
        item = {
            "id": row.account_id,
            "account_id": row.account_id,
            "account_name": row.account_name,
            "user_id": row.primary_clerk_user_id or row.primary_user_id,
            "primary_email": primary_email_value,
            "email_domain": _email_domain(primary_email_value),
            "identity": identity,
            "surface_flags": surface_flags,
            "surface_reasons": signals.reasons,
            "active_surfaces": active_surfaces,
            "primary_surface": primary_surface,
            "actions_blocked": bool(actions_blocked_reason),
            "actions_blocked_reason": actions_blocked_reason,
            "license_type": "paid" if resolved_tier != "free" else "trial",
            "tier": resolved_tier,
            "status": status,
            "status_detail": (
                "Account exists, but no trial quota or entitlement row has been provisioned yet."
                if orphan_account
                else None
            ),
            "record_kind": "orphan_account" if orphan_account else "trial_account",
            "trial_started_at": _iso_or_none(row.trial_started_at),
            "trial_ends_at": _iso_or_none(row.expires_at),
            "days_remaining": _days_remaining(row.expires_at),
            "sync_source": row.tier_sync_source or row.entitlement_source,
            "updated_at": _iso_or_none(row.tier_synced_at or row.primary_user_created_at),
            "submitted_at": None,
            "invited_at": None,
        }
        items.append(item)
        if primary_email_value:
            surfaces_for_dedupe = list(active_surfaces) or [primary_surface]
            for surface in surfaces_for_dedupe:
                accounted_email_surface_pairs.add((primary_email_value, surface))

    beta_result = await db.execute(
        select(BetaSubmission).where(BetaSubmission.status.in_(("pending", "sent")))
    )
    lead_items: dict[tuple[str, str], Dict[str, Any]] = {}
    for submission in beta_result.scalars().all():
        email = (submission.email or "").strip().lower()
        if not email:
            continue
        surface = normalize_product_surface(submission.product_surface)
        if (email, surface) in accounted_email_surface_pairs:
            continue

        status = "invited" if (submission.status or "").strip().lower() == "sent" else "pending"
        lead_items[(email, surface)] = {
            "id": f"beta:{submission.id}",
            "account_id": None,
            "account_name": submission.domain or submission.name or email,
            "user_id": None,
            "primary_email": email,
            "email_domain": submission.domain or _email_domain(email),
            "identity": _lead_identity(
                label="Beta signup",
                reason="Lead has not activated or provisioned an account yet",
            ),
            "surface_flags": _surface_flags_for(surface),
            "surface_reasons": _surface_reasons_for(
                surface,
                f"beta_submission:{submission.status}:{submission.source}",
            ),
            "active_surfaces": [surface],
            "primary_surface": surface,
            "actions_blocked": True,
            "actions_blocked_reason": (
                "No account has signed in yet. This row is still a beta lead or invitation, not an activated account."
            ),
            "license_type": "invite" if status == "invited" else "beta",
            "tier": "free",
            "status": status,
            "status_detail": (
                "Invitation sent, but the user has not activated a trial account yet."
                if status == "invited"
                else "Beta signup captured, but no invitation or activated account exists yet."
            ),
            "record_kind": "beta_submission",
            "trial_started_at": None,
            "trial_ends_at": None,
            "days_remaining": None,
            "sync_source": submission.source,
            "updated_at": _iso_or_none(submission.reviewed_at or submission.submitted_at),
            "submitted_at": _iso_or_none(submission.submitted_at),
            "invited_at": _iso_or_none(submission.reviewed_at) if status == "invited" else None,
        }

    invite_result = await db.execute(
        select(
            TrialEmailSend.recipient_email,
            TrialEmailSend.sent_at,
            TrialEmailTemplate.product_surface,
        )
        .join(TrialEmailTemplate, TrialEmailTemplate.id == TrialEmailSend.template_id)
        .where(
            TrialEmailTemplate.slug == "manual_invitation",
            TrialEmailSend.status == "sent",
        )
    )
    for recipient_email, sent_at, product_surface in invite_result.all():
        email = (recipient_email or "").strip().lower()
        if not email:
            continue
        surface = normalize_product_surface(product_surface)
        if (email, surface) in accounted_email_surface_pairs:
            continue

        lead_key = (email, surface)
        invited_at = _iso_or_none(sent_at)
        if lead_key in lead_items:
            lead_item = lead_items[lead_key]
            lead_item["status"] = "invited"
            lead_item["license_type"] = "invite"
            lead_item["invited_at"] = invited_at or lead_item.get("invited_at")
            lead_item["updated_at"] = invited_at or lead_item.get("updated_at")
            lead_item["status_detail"] = (
                "Invitation sent, but the user has not activated a trial account yet."
            )
            lead_item["surface_reasons"][surface].append("trial_email:manual_invitation")
            continue

        lead_items[lead_key] = {
            "id": f"invite:{surface}:{email}",
            "account_id": None,
            "account_name": _email_domain(email) or email,
            "user_id": None,
            "primary_email": email,
            "email_domain": _email_domain(email),
            "identity": _lead_identity(
                label="Invited lead",
                reason="Invitation sent before account activation",
            ),
            "surface_flags": _surface_flags_for(surface),
            "surface_reasons": _surface_reasons_for(
                surface,
                "trial_email:manual_invitation",
            ),
            "active_surfaces": [surface],
            "primary_surface": surface,
            "actions_blocked": True,
            "actions_blocked_reason": (
                "No account has signed in yet. This row only reflects an invitation email."
            ),
            "license_type": "invite",
            "tier": "free",
            "status": "invited",
            "status_detail": "Invitation sent, but no activated account exists yet.",
            "record_kind": "manual_invitation",
            "trial_started_at": None,
            "trial_ends_at": None,
            "days_remaining": None,
            "sync_source": "manual_invitation",
            "updated_at": invited_at,
            "submitted_at": None,
            "invited_at": invited_at,
        }

    items.extend(lead_items.values())
    items = [item for item in items if _search_matches(item, search)]
    items.sort(key=_inventory_sort_key)

    return items


async def _load_trial_account(
    db: AsyncSession,
    *,
    account_id: str,
) -> tuple[Account, Optional[UsageQuotaModel], Optional[AccountEntitlements]]:
    account = await db.get(Account, account_id)
    if not account:
        raise HTTPException(404, "Account not found")

    quota = await db.get(UsageQuotaModel, account_id)
    entitlements = await db.get(AccountEntitlements, account_id)
    return account, quota, entitlements


async def _load_trial_inventory_item(
    db: AsyncSession,
    *,
    account_id: str,
) -> Dict[str, Any]:
    items = await _fetch_trial_inventory(db, search=account_id)
    for item in items:
        if item["account_id"] == account_id:
            return item
    raise HTTPException(404, "Trial account not found")


def _resolve_requested_surface(
    *,
    trial_item: Dict[str, Any],
    requested_surface: Optional[str],
) -> str:
    active_surfaces = list(trial_item.get("active_surfaces") or [])
    blocked_reason = trial_item.get("actions_blocked_reason")
    if blocked_reason:
        raise HTTPException(409, blocked_reason)

    if requested_surface:
        normalized_requested = normalize_product_surface(requested_surface)
        if active_surfaces and normalized_requested not in active_surfaces:
            raise HTTPException(
                400,
                f"Account is not currently attributed to the {normalized_requested} surface",
            )
        return normalized_requested

    primary_surface = normalize_product_surface(
        trial_item.get("primary_surface"),
        default=DEFAULT_PRODUCT_SURFACE,
    )
    if primary_surface in active_surfaces:
        return primary_surface
    if len(active_surfaces) == 1:
        return active_surfaces[0]
    return primary_surface


@router.get("")
async def list_trials(
    admin_ctx: AdminContext = Depends(require_admin),
    db: AsyncSession = Depends(get_session),
    status: str = Query(default="all"),
    search: Optional[str] = Query(default=None),
) -> Dict[str, Any]:
    try:
        items = await _fetch_trial_inventory(db, search=search)
        normalized_status = status.strip().lower()
        if normalized_status not in TRIAL_STATUS_FILTERS:
            raise HTTPException(400, "Unsupported trial status filter")
        if normalized_status != "all":
            items = [item for item in items if item["status"] == normalized_status]

        await log_admin_action(
            admin_user_id=admin_ctx.user_id,
            admin_email=admin_ctx.email,
            action="list_trials",
            db=db,
            details={"status": normalized_status, "search": search},
        )
        return {"ok": True, "trials": items, "total": len(items)}
    except HTTPException:
        raise
    except Exception as exc:
        log.error("Failed to list trials: %s", exc, exc_info=True)
        raise HTTPException(500, f"Failed to list trials: {exc}")


@router.post("/{account_id}/extend")
async def extend_trial(
    account_id: str,
    payload: TrialExtendPayload = Body(...),
    admin_ctx: AdminContext = Depends(require_admin),
    db: AsyncSession = Depends(get_session),
) -> Dict[str, Any]:
    account, quota, entitlements = await _load_trial_account(db, account_id=account_id)
    trial_item = await _load_trial_inventory_item(db, account_id=account_id)
    product_surface = _resolve_requested_surface(
        trial_item=trial_item,
        requested_surface=payload.product_surface,
    )
    action_source = _build_trial_action_source(
        action="extend",
        product_surface=product_surface,
    )
    account_name = account.name
    current_tier = (
        (entitlements.tier if entitlements else None)
        or (quota.tier if quota else None)
        or "free"
    ).strip().lower()
    if current_tier != "free":
        raise HTTPException(400, "Only free-tier trials can be extended")

    monthly_rows, monthly_enrichments = _resolve_entitlements_from_limits(
        "free", None, None
    )
    quota = await _apply_tier_update(
        account_id=account_id,
        tier="free",
        status="trialing",
        monthly_rows=monthly_rows,
        monthly_enrichments=monthly_enrichments,
        source=action_source,
        addons=[],
        db=db,
    )
    response_tier = quota.tier
    response_status = quota.subscription_status

    now = datetime.now(timezone.utc)
    existing_entitlements = await db.get(AccountEntitlements, account_id)
    current_expiry = _ensure_utc(existing_entitlements.expires_at) if existing_entitlements else None
    base = current_expiry if current_expiry and current_expiry > now else now
    entitlements = await _upsert_trial_entitlements(
        db,
        account_id=account_id,
        tier="free",
        source=action_source,
        monthly_rows=monthly_rows,
        trial_started_at=(existing_entitlements.trial_started_at if existing_entitlements and existing_entitlements.trial_started_at else now),
        expires_at=base + timedelta(days=payload.extend_days),
    )

    await log_admin_action(
        admin_user_id=admin_ctx.user_id,
        admin_email=admin_ctx.email,
        action="extend_trial",
        db=db,
        target_account_id=account_id,
        details={
            "account_name": account_name,
            "product_surface": product_surface,
            "extend_days": payload.extend_days,
            "new_trial_end": _iso_or_none(entitlements.expires_at),
        },
    )

    return {
        "ok": True,
        "account_id": account_id,
        "product_surface": product_surface,
        "tier": response_tier,
        "status": response_status,
        "trial_ends_at": _iso_or_none(entitlements.expires_at),
    }


@router.post("/{account_id}/revoke")
async def revoke_trial(
    account_id: str,
    payload: TrialRevokePayload = Body(default=TrialRevokePayload()),
    admin_ctx: AdminContext = Depends(require_admin),
    db: AsyncSession = Depends(get_session),
) -> Dict[str, Any]:
    account, _quota, _entitlements = await _load_trial_account(db, account_id=account_id)
    trial_item = await _load_trial_inventory_item(db, account_id=account_id)
    product_surface = _resolve_requested_surface(
        trial_item=trial_item,
        requested_surface=payload.product_surface,
    )
    action_source = _build_trial_action_source(
        action="revoke",
        product_surface=product_surface,
    )
    account_name = account.name
    monthly_rows, monthly_enrichments = _resolve_entitlements_from_limits(
        "free", None, None
    )
    quota = await _apply_tier_update(
        account_id=account_id,
        tier="free",
        status="revoked",
        monthly_rows=monthly_rows,
        monthly_enrichments=monthly_enrichments,
        source=action_source,
        addons=[],
        db=db,
    )
    response_tier = quota.tier
    response_status = quota.subscription_status

    now = datetime.now(timezone.utc)
    existing_entitlements = await db.get(AccountEntitlements, account_id)
    await _upsert_trial_entitlements(
        db,
        account_id=account_id,
        tier="free",
        source=action_source,
        monthly_rows=monthly_rows,
        trial_started_at=(
            existing_entitlements.trial_started_at
            if existing_entitlements and existing_entitlements.trial_started_at
            else now
        ),
        expires_at=now,
    )

    await log_admin_action(
        admin_user_id=admin_ctx.user_id,
        admin_email=admin_ctx.email,
        action="revoke_trial",
        db=db,
        target_account_id=account_id,
        details={
            "account_name": account_name,
            "product_surface": product_surface,
            "reason": payload.reason,
        },
    )

    return {
        "ok": True,
        "account_id": account_id,
        "product_surface": product_surface,
        "tier": response_tier,
        "status": response_status,
        "revoked_at": datetime.now(timezone.utc).isoformat(),
    }


@router.post("/{account_id}/upgrade")
async def upgrade_trial(
    account_id: str,
    payload: TrialUpgradePayload = Body(default=TrialUpgradePayload()),
    admin_ctx: AdminContext = Depends(require_admin),
    db: AsyncSession = Depends(get_session),
) -> Dict[str, Any]:
    account, _quota, _entitlements = await _load_trial_account(db, account_id=account_id)
    trial_item = await _load_trial_inventory_item(db, account_id=account_id)
    product_surface = _resolve_requested_surface(
        trial_item=trial_item,
        requested_surface=payload.product_surface,
    )
    action_source = _build_trial_action_source(
        action="upgrade",
        product_surface=product_surface,
    )
    account_name = account.name
    monthly_rows, monthly_enrichments = _resolve_entitlements_from_limits(
        payload.tier, None, None
    )
    quota = await _apply_tier_update(
        account_id=account_id,
        tier=payload.tier,
        status="active",
        monthly_rows=monthly_rows,
        monthly_enrichments=monthly_enrichments,
        source=action_source,
        addons=[],
        db=db,
    )
    response_tier = quota.tier
    response_status = quota.subscription_status
    response_source = quota.tier_sync_source

    existing_entitlements = await db.get(AccountEntitlements, account_id)
    await _upsert_trial_entitlements(
        db,
        account_id=account_id,
        tier=payload.tier,
        source=action_source,
        monthly_rows=monthly_rows,
        trial_started_at=(
            existing_entitlements.trial_started_at
            if existing_entitlements
            else datetime.now(timezone.utc)
        ),
        expires_at=None,
    )

    await log_admin_action(
        admin_user_id=admin_ctx.user_id,
        admin_email=admin_ctx.email,
        action="upgrade_trial",
        db=db,
        target_account_id=account_id,
        details={
            "account_name": account_name,
            "tier": payload.tier,
            "product_surface": product_surface,
        },
    )

    return {
        "ok": True,
        "account_id": account_id,
        "product_surface": product_surface,
        "tier": response_tier,
        "status": response_status,
        "source": response_source,
    }
