"""Billing subscription model - polymorphic subscriptions for any resource type."""

from datetime import UTC, datetime, timedelta
from enum import Enum
from typing import Optional

from sqlalchemy import Column, DateTime, ForeignKey, Index, Integer, String
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session

from robosystems.database import Base
from robosystems.logger import get_logger
from robosystems.utils.ulid import generate_prefixed_ulid

logger = get_logger(__name__)


class SubscriptionStatus(str, Enum):
  """Subscription status states."""

  PENDING = "pending"
  PENDING_PAYMENT = "pending_payment"
  PROVISIONING = "provisioning"
  ACTIVE = "active"
  UPGRADING = "upgrading"
  PAUSED = "paused"
  CANCELED = "canceled"
  PAST_DUE = "past_due"
  UNPAID = "unpaid"


class BillingInterval(str, Enum):
  """Billing interval options."""

  MONTHLY = "monthly"
  ANNUAL = "annual"
  USAGE_BASED = "usage_based"


class CancellationType(str, Enum):
  """How a subscription was canceled."""

  PERIOD_END = "period_end"
  IMMEDIATE = "immediate"


class BillingSubscription(Base):
  """Generic subscription model for any billable resource.

  Designed to be polymorphic - can bill for graphs, repositories,
  API add-ons, storage, or any future billable resource.

  Separated from resource models to isolate billing concerns.
  """

  __tablename__ = "billing_subscriptions"

  id = Column(String, primary_key=True, default=lambda: generate_prefixed_ulid("bsub"))

  org_id = Column(String, ForeignKey("orgs.id"), nullable=False)

  resource_type = Column(String, nullable=False)
  resource_id = Column(String, nullable=True)

  plan_name = Column(String, nullable=False)
  billing_interval = Column(String, default="monthly", nullable=False)

  base_price_cents = Column(Integer, nullable=False)

  stripe_subscription_id = Column(String, unique=True, nullable=True)
  stripe_product_id = Column(String, nullable=True)
  stripe_price_id = Column(String, nullable=True)

  payment_provider = Column(String, default="invoice", nullable=False)
  provider_subscription_id = Column(String, unique=True, nullable=True)
  provider_customer_id = Column(String, nullable=True)

  subscription_metadata = Column(
    JSONB, default=dict, nullable=False, server_default="{}"
  )

  status = Column(String, default="pending", nullable=False)
  cancellation_type = Column(String, nullable=True)

  started_at = Column(DateTime, nullable=True)
  current_period_start = Column(DateTime, nullable=True)
  current_period_end = Column(DateTime, nullable=True)
  canceled_at = Column(DateTime, nullable=True)
  ends_at = Column(DateTime, nullable=True)

  created_at = Column(DateTime, default=lambda: datetime.now(UTC), nullable=False)
  updated_at = Column(
    DateTime,
    default=lambda: datetime.now(UTC),
    onupdate=lambda: datetime.now(UTC),
    nullable=False,
  )

  __table_args__ = (
    Index("idx_billing_sub_org", "org_id"),
    Index("idx_billing_sub_resource", "resource_type", "resource_id"),
    Index("idx_billing_sub_status", "status"),
    Index("idx_billing_sub_stripe", "stripe_subscription_id"),
    Index("idx_billing_sub_provider", "provider_subscription_id"),
    Index("idx_billing_sub_cancellation_type", "cancellation_type"),
  )

  def __repr__(self) -> str:
    return f"<BillingSubscription {self.resource_type}:{self.resource_id} plan={self.plan_name}>"

  @classmethod
  def create_subscription(
    cls,
    org_id: str,
    resource_type: str,
    resource_id: str,
    plan_name: str,
    base_price_cents: int,
    session: Session,
    billing_interval: str = "monthly",
    stripe_subscription_id: str | None = None,
  ) -> "BillingSubscription":
    """Create a new subscription."""
    now = datetime.now(UTC)

    subscription = cls(
      org_id=org_id,
      resource_type=resource_type,
      resource_id=resource_id,
      plan_name=plan_name,
      base_price_cents=base_price_cents,
      billing_interval=billing_interval,
      stripe_subscription_id=stripe_subscription_id,
      status=SubscriptionStatus.PENDING.value,
      created_at=now,
    )

    session.add(subscription)
    session.commit()
    session.refresh(subscription)

    logger.info(
      f"Created billing subscription {subscription.id} for {resource_type}:{resource_id}"
    )

    return subscription

  @classmethod
  def get_by_resource(
    cls, resource_type: str, resource_id: str, session: Session
  ) -> Optional["BillingSubscription"]:
    """Get subscription for a specific resource."""
    return (
      session.query(cls)
      .filter(cls.resource_type == resource_type, cls.resource_id == resource_id)
      .first()
    )

  @classmethod
  def get_by_resource_and_org(
    cls,
    resource_type: str,
    resource_id: str,
    org_id: str,
    session: Session,
  ) -> Optional["BillingSubscription"]:
    """Get subscription for a specific resource and organization.

    This is particularly useful for shared repositories where multiple orgs
    can have separate subscriptions to the same resource.
    """
    return (
      session.query(cls)
      .filter(
        cls.resource_type == resource_type,
        cls.resource_id == resource_id,
        cls.org_id == org_id,
      )
      .first()
    )

  @classmethod
  def get_by_resource_and_user(
    cls,
    resource_type: str,
    resource_id: str,
    user_id: str,
    session: Session,
  ) -> Optional["BillingSubscription"]:
    """Get subscription for a specific resource and user (looks up user's org first)."""
    from robosystems.models.core.org import OrgUser

    membership = session.query(OrgUser).filter(OrgUser.user_id == user_id).first()

    if not membership:
      return None

    return cls.get_by_resource_and_org(
      resource_type=resource_type,
      resource_id=resource_id,
      org_id=membership.org_id,
      session=session,
    )

  @classmethod
  def get_active_subscriptions_for_org(
    cls, org_id: str, session: Session
  ) -> list["BillingSubscription"]:
    """Get all active subscriptions for an organization."""
    return (
      session.query(cls)
      .filter(
        cls.org_id == org_id,
        cls.status == SubscriptionStatus.ACTIVE.value,
      )
      .all()
    )

  @classmethod
  def get_active_subscriptions_for_user(
    cls, user_id: str, session: Session
  ) -> list["BillingSubscription"]:
    """Get all active subscriptions for a user (looks up user's org first)."""
    from robosystems.models.core.org import OrgUser

    membership = session.query(OrgUser).filter(OrgUser.user_id == user_id).first()

    if not membership:
      return []

    return cls.get_active_subscriptions_for_org(membership.org_id, session)

  @classmethod
  def get_by_provider_subscription_id(
    cls, provider_subscription_id: str, session: Session
  ) -> Optional["BillingSubscription"]:
    """Get subscription by payment provider subscription ID (e.g., Stripe subscription ID)."""
    return (
      session.query(cls)
      .filter(cls.provider_subscription_id == provider_subscription_id)
      .first()
    )

  @classmethod
  def get_by_stripe_subscription_id(
    cls, stripe_subscription_id: str, session: Session
  ) -> Optional["BillingSubscription"]:
    """Get subscription by Stripe subscription ID (legacy support)."""
    return (
      session.query(cls)
      .filter(cls.stripe_subscription_id == stripe_subscription_id)
      .first()
    )

  def _invalidate_access_cache(self) -> None:
    """Invalidate the subscription access cache for this resource.

    Called on any status change so require_graph_access() re-checks the DB
    instead of serving a stale cached result for up to 300s.
    """
    if self.resource_type == "graph" and self.resource_id:
      try:
        from robosystems.middleware.billing.enforcement import (
          invalidate_subscription_cache,
        )

        invalidate_subscription_cache(self.resource_id)
      except Exception as e:
        logger.warning(f"Cache invalidation failed for graph {self.resource_id}: {e}")

  def activate(self, session: Session) -> None:
    """Activate the subscription."""
    now = datetime.now(UTC)
    self.status = SubscriptionStatus.ACTIVE.value
    self.started_at = now
    self.current_period_start = now
    self.current_period_end = now + timedelta(days=30)
    self.updated_at = now

    session.commit()
    session.refresh(self)
    self._invalidate_access_cache()

    logger.info(f"Activated subscription {self.id}")

  def pause(self, session: Session) -> None:
    """Pause the subscription."""
    self.status = SubscriptionStatus.PAUSED.value
    self.updated_at = datetime.now(UTC)

    session.commit()
    session.refresh(self)
    self._invalidate_access_cache()

    logger.info(f"Paused subscription {self.id}")

  def cancel(self, session: Session, immediate: bool = False) -> None:
    """Cancel the subscription.

    Sets `cancellation_type` to "immediate" or "period_end" so downstream
    sensors (e.g. deprovisioning) can branch on the user's intent without
    re-deriving it from timestamps.

    Raises ValueError if `immediate=False` is requested on a subscription
    whose `current_period_end` is None — without an `ends_at` value the
    deprovision sensor's `ends_at.isnot(None)` filter would silently skip
    the sub forever, leaking infrastructure.
    """
    now = datetime.now(UTC)

    if not immediate and self.current_period_end is None:
      raise ValueError(
        f"Cannot cancel subscription {self.id} at period end: "
        "current_period_end is None (subscription was never activated). "
        "Use immediate=True instead."
      )

    self.status = SubscriptionStatus.CANCELED.value
    self.canceled_at = now

    if immediate:
      self.ends_at = now
      self.cancellation_type = CancellationType.IMMEDIATE.value
    else:
      self.ends_at = self.current_period_end
      self.cancellation_type = CancellationType.PERIOD_END.value

    self.updated_at = now

    session.commit()
    session.refresh(self)
    self._invalidate_access_cache()

    logger.info(
      f"Canceled subscription {self.id} "
      f"(type={self.cancellation_type}, ends={self.ends_at})"
    )

  def update_plan(
    self, new_plan_name: str, new_price_cents: int, session: Session
  ) -> None:
    """Update subscription plan (upgrade or downgrade)."""
    old_plan = self.plan_name
    self.plan_name = new_plan_name
    self.base_price_cents = new_price_cents
    self.updated_at = datetime.now(UTC)

    session.commit()
    session.refresh(self)
    self._invalidate_access_cache()

    logger.info(f"Updated subscription {self.id} plan: {old_plan} -> {new_plan_name}")

  def update_stripe_subscription(
    self,
    stripe_subscription_id: str,
    stripe_product_id: str,
    stripe_price_id: str,
    session: Session,
  ) -> None:
    """Update Stripe subscription IDs."""
    self.stripe_subscription_id = stripe_subscription_id
    self.stripe_product_id = stripe_product_id
    self.stripe_price_id = stripe_price_id
    self.updated_at = datetime.now(UTC)

    session.commit()
    session.refresh(self)

    logger.info(f"Updated Stripe subscription for {self.id}")

  def _get_period_delta(self) -> timedelta:
    """Get the timedelta for one billing period based on billing_interval."""
    if self.billing_interval == BillingInterval.ANNUAL.value:
      return timedelta(days=365)
    # Monthly is the default for both "monthly" and "usage_based"
    return timedelta(days=30)

  def renew_period(self, session: Session) -> None:
    """Advance to the next billing period.

    Shifts current_period_start to the old current_period_end and extends
    current_period_end by the billing interval duration. Used by the invoice
    billing renewal job.
    """
    now = datetime.now(UTC)
    delta = self._get_period_delta()
    self.current_period_start = self.current_period_end
    self.current_period_end = self.current_period_end + delta
    self.updated_at = now

    session.flush()

    logger.info(f"Renewed subscription {self.id} period to {self.current_period_end}")

  def is_active(self) -> bool:
    """Check if subscription is currently active."""
    return self.status == SubscriptionStatus.ACTIVE.value
