"""
Subscription and Usage Tracking Models
Comprehensive models for usage-based subscription system with API cost tracking.
"""

from sqlalchemy import Column, Integer, String, DateTime, Float, Boolean, JSON, Text, ForeignKey, Enum
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship
from datetime import datetime, timedelta
import enum
from typing import Dict, Any, Optional

Base = declarative_base()

class SubscriptionTier(enum.Enum):
    FREE = "free"
    BASIC = "basic"
    PRO = "pro"
    ENTERPRISE = "enterprise"

class UsageStatus(enum.Enum):
    ACTIVE = "active"
    WARNING = "warning"  # 80% usage
    LIMIT_REACHED = "limit_reached"  # 100% usage
    SUSPENDED = "suspended"

class APIProvider(enum.Enum):
    GEMINI = "gemini"
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    MISTRAL = "mistral"
    TAVILY = "tavily"
    SERPER = "serper"
    METAPHOR = "metaphor"
    FIRECRAWL = "firecrawl"
    STABILITY = "stability"
    EXA = "exa"
    VIDEO = "video"
    IMAGE_EDIT = "image_edit"
    AUDIO = "audio"

class BillingCycle(enum.Enum):
    MONTHLY = "monthly"
    YEARLY = "yearly"

class SubscriptionPlan(Base):
    """Defines subscription tiers and their limits."""
    
    __tablename__ = "subscription_plans"
    
    id = Column(Integer, primary_key=True)
    name = Column(String(50), nullable=False, unique=True)
    tier = Column(Enum(SubscriptionTier), nullable=False)
    price_monthly = Column(Float, nullable=False, default=0.0)
    price_yearly = Column(Float, nullable=False, default=0.0)
    
    # Unified AI Text Generation Call Limit (applies to all LLM providers: gemini, openai, anthropic, mistral)
    # Note: This column may not exist in older databases - use getattr() when accessing
    ai_text_generation_calls_limit = Column(Integer, default=0, nullable=True)  # 0 = unlimited, None if column doesn't exist
    
    # Legacy per-provider limits (kept for backwards compatibility and analytics)
    gemini_calls_limit = Column(Integer, default=0)  # 0 = unlimited (deprecated, use ai_text_generation_calls_limit)
    openai_calls_limit = Column(Integer, default=0)  # (deprecated, use ai_text_generation_calls_limit)
    anthropic_calls_limit = Column(Integer, default=0)  # (deprecated, use ai_text_generation_calls_limit)
    mistral_calls_limit = Column(Integer, default=0)  # (deprecated, use ai_text_generation_calls_limit)
    
    # Other API Call Limits (non-LLM)
    tavily_calls_limit = Column(Integer, default=0)
    serper_calls_limit = Column(Integer, default=0)
    metaphor_calls_limit = Column(Integer, default=0)
    firecrawl_calls_limit = Column(Integer, default=0)
    stability_calls_limit = Column(Integer, default=0)  # Image generation
    exa_calls_limit = Column(Integer, default=0)  # Exa neural search
    video_calls_limit = Column(Integer, default=0)  # AI video generation
    image_edit_calls_limit = Column(Integer, default=0)  # AI image editing
    audio_calls_limit = Column(Integer, default=0)  # AI audio generation (text-to-speech)
    
    # Token Limits (for LLM providers)
    gemini_tokens_limit = Column(Integer, default=0)
    openai_tokens_limit = Column(Integer, default=0)
    anthropic_tokens_limit = Column(Integer, default=0)
    mistral_tokens_limit = Column(Integer, default=0)
    
    # Cost Limits (in USD)
    monthly_cost_limit = Column(Float, default=0.0)  # 0 = unlimited
    
    # Features
    features = Column(JSON, nullable=True)  # JSON list of enabled features
    
    # Metadata
    description = Column(Text, nullable=True)
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class UserSubscription(Base):
    """User's current subscription and billing information."""
    
    __tablename__ = "user_subscriptions"
    
    id = Column(Integer, primary_key=True)
    user_id = Column(String(100), nullable=False, unique=True)
    plan_id = Column(Integer, ForeignKey('subscription_plans.id'), nullable=False)
    
    # Billing
    billing_cycle = Column(Enum(BillingCycle), default=BillingCycle.MONTHLY)
    current_period_start = Column(DateTime, nullable=False)
    current_period_end = Column(DateTime, nullable=False)
    
    # Status
    status = Column(Enum(UsageStatus), default=UsageStatus.ACTIVE)
    is_active = Column(Boolean, default=True)
    auto_renew = Column(Boolean, default=True)
    
    # Payment
    stripe_customer_id = Column(String(100), nullable=True)
    stripe_subscription_id = Column(String(100), nullable=True)
    payment_method = Column(String(50), nullable=True)
    
    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # Relationships
    plan = relationship("SubscriptionPlan")

class APIUsageLog(Base):
    """Detailed log of every API call for billing and monitoring."""
    
    __tablename__ = "api_usage_logs"
    
    id = Column(Integer, primary_key=True)
    user_id = Column(String(100), nullable=False)
    
    # API Details
    provider = Column(Enum(APIProvider), nullable=False)
    endpoint = Column(String(200), nullable=False)
    method = Column(String(10), nullable=False)
    model_used = Column(String(100), nullable=True)  # e.g., "gemini-2.5-flash"
    
    # Usage Metrics
    tokens_input = Column(Integer, default=0)
    tokens_output = Column(Integer, default=0)
    tokens_total = Column(Integer, default=0)
    
    # Cost Calculation
    cost_input = Column(Float, default=0.0)  # Cost for input tokens
    cost_output = Column(Float, default=0.0)  # Cost for output tokens
    cost_total = Column(Float, default=0.0)  # Total cost for this call
    
    # Performance
    response_time = Column(Float, nullable=False)  # Response time in seconds
    status_code = Column(Integer, nullable=False)
    
    # Request Details
    request_size = Column(Integer, nullable=True)  # Request size in bytes
    response_size = Column(Integer, nullable=True)  # Response size in bytes
    user_agent = Column(String(500), nullable=True)
    ip_address = Column(String(45), nullable=True)
    
    # Error Handling
    error_message = Column(Text, nullable=True)
    retry_count = Column(Integer, default=0)
    
    # Metadata
    timestamp = Column(DateTime, default=datetime.utcnow, nullable=False)
    billing_period = Column(String(20), nullable=False)  # e.g., "2025-01"
    
    # Indexes for performance
    __table_args__ = (
        {'mysql_engine': 'InnoDB'},
    )

class UsageSummary(Base):
    """Aggregated usage statistics per user per billing period."""
    
    __tablename__ = "usage_summaries"
    
    id = Column(Integer, primary_key=True)
    user_id = Column(String(100), nullable=False)
    billing_period = Column(String(20), nullable=False)  # e.g., "2025-01"
    
    # API Call Counts
    gemini_calls = Column(Integer, default=0)
    openai_calls = Column(Integer, default=0)
    anthropic_calls = Column(Integer, default=0)
    mistral_calls = Column(Integer, default=0)
    tavily_calls = Column(Integer, default=0)
    serper_calls = Column(Integer, default=0)
    metaphor_calls = Column(Integer, default=0)
    firecrawl_calls = Column(Integer, default=0)
    stability_calls = Column(Integer, default=0)
    exa_calls = Column(Integer, default=0)
    video_calls = Column(Integer, default=0)  # AI video generation
    image_edit_calls = Column(Integer, default=0)  # AI image editing
    audio_calls = Column(Integer, default=0)  # AI audio generation (text-to-speech)
    
    # Token Usage
    gemini_tokens = Column(Integer, default=0)
    openai_tokens = Column(Integer, default=0)
    anthropic_tokens = Column(Integer, default=0)
    mistral_tokens = Column(Integer, default=0)
    
    # Cost Tracking
    gemini_cost = Column(Float, default=0.0)
    openai_cost = Column(Float, default=0.0)
    anthropic_cost = Column(Float, default=0.0)
    mistral_cost = Column(Float, default=0.0)
    tavily_cost = Column(Float, default=0.0)
    serper_cost = Column(Float, default=0.0)
    metaphor_cost = Column(Float, default=0.0)
    firecrawl_cost = Column(Float, default=0.0)
    stability_cost = Column(Float, default=0.0)
    exa_cost = Column(Float, default=0.0)
    video_cost = Column(Float, default=0.0)  # AI video generation
    image_edit_cost = Column(Float, default=0.0)  # AI image editing
    audio_cost = Column(Float, default=0.0)  # AI audio generation (text-to-speech)
    
    # Totals
    total_calls = Column(Integer, default=0)
    total_tokens = Column(Integer, default=0)
    total_cost = Column(Float, default=0.0)
    
    # Performance Metrics
    avg_response_time = Column(Float, default=0.0)
    error_rate = Column(Float, default=0.0)  # Percentage
    
    # Status
    usage_status = Column(Enum(UsageStatus), default=UsageStatus.ACTIVE)
    warnings_sent = Column(Integer, default=0)  # Number of warning emails sent
    
    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # Unique constraint on user_id and billing_period
    __table_args__ = (
        {'mysql_engine': 'InnoDB'},
    )

class APIProviderPricing(Base):
    """Pricing configuration for different API providers."""
    
    __tablename__ = "api_provider_pricing"
    
    id = Column(Integer, primary_key=True)
    provider = Column(Enum(APIProvider), nullable=False)
    model_name = Column(String(100), nullable=False)
    
    # Pricing per token (in USD)
    cost_per_input_token = Column(Float, default=0.0)
    cost_per_output_token = Column(Float, default=0.0)
    cost_per_request = Column(Float, default=0.0)  # Fixed cost per API call
    
    # Pricing per unit for non-LLM APIs
    cost_per_search = Column(Float, default=0.0)  # For search APIs
    cost_per_image = Column(Float, default=0.0)  # For image generation
    cost_per_page = Column(Float, default=0.0)  # For web crawling
    
    # Token conversion (tokens per unit)
    tokens_per_word = Column(Float, default=1.3)  # Approximate tokens per word
    tokens_per_character = Column(Float, default=0.25)  # Approximate tokens per character
    
    # Metadata
    description = Column(Text, nullable=True)
    is_active = Column(Boolean, default=True)
    effective_date = Column(DateTime, default=datetime.utcnow)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # Unique constraint on provider and model
    __table_args__ = (
        {'mysql_engine': 'InnoDB'},
    )

class UsageAlert(Base):
    """Usage alerts and notifications."""
    
    __tablename__ = "usage_alerts"
    
    id = Column(Integer, primary_key=True)
    user_id = Column(String(100), nullable=False)
    
    # Alert Details
    alert_type = Column(String(50), nullable=False)  # "usage_warning", "limit_reached", "cost_warning"
    threshold_percentage = Column(Integer, nullable=False)  # 80, 90, 100
    provider = Column(Enum(APIProvider), nullable=True)  # Specific provider or None for overall
    
    # Alert Content
    title = Column(String(200), nullable=False)
    message = Column(Text, nullable=False)
    severity = Column(String(20), default="info")  # "info", "warning", "error"
    
    # Status
    is_sent = Column(Boolean, default=False)
    sent_at = Column(DateTime, nullable=True)
    is_read = Column(Boolean, default=False)
    read_at = Column(DateTime, nullable=True)
    
    # Metadata
    billing_period = Column(String(20), nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    
class BillingHistory(Base):
    """Historical billing records."""
    
    __tablename__ = "billing_history"
    
    id = Column(Integer, primary_key=True)
    user_id = Column(String(100), nullable=False)
    
    # Billing Period
    billing_period = Column(String(20), nullable=False)  # e.g., "2025-01"
    period_start = Column(DateTime, nullable=False)
    period_end = Column(DateTime, nullable=False)
    
    # Subscription
    plan_name = Column(String(50), nullable=False)
    base_cost = Column(Float, default=0.0)
    
    # Usage Costs
    usage_cost = Column(Float, default=0.0)
    overage_cost = Column(Float, default=0.0)
    total_cost = Column(Float, default=0.0)
    
    # Payment
    payment_status = Column(String(20), default="pending")  # "pending", "paid", "failed"
    payment_date = Column(DateTime, nullable=True)
    stripe_invoice_id = Column(String(100), nullable=True)
    
    # Usage Summary (snapshot)
    total_api_calls = Column(Integer, default=0)
    total_tokens = Column(Integer, default=0)
    usage_details = Column(JSON, nullable=True)  # Detailed breakdown
    
    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class SubscriptionRenewalHistory(Base):
    """Historical record of subscription renewals and expiration events."""
    
    __tablename__ = "subscription_renewal_history"
    
    id = Column(Integer, primary_key=True)
    user_id = Column(String(100), nullable=False)
    
    # Subscription Details
    plan_id = Column(Integer, ForeignKey('subscription_plans.id'), nullable=False)
    plan_name = Column(String(50), nullable=False)
    plan_tier = Column(String(20), nullable=False)  # e.g., "free", "basic", "pro", "enterprise"
    
    # Period Information
    previous_period_start = Column(DateTime, nullable=True)  # Start of the previous period (if renewal)
    previous_period_end = Column(DateTime, nullable=True)  # End of the previous period (when it expired)
    new_period_start = Column(DateTime, nullable=False)  # Start of the new period (when renewed)
    new_period_end = Column(DateTime, nullable=False)  # End of the new period
    
    # Billing Cycle
    billing_cycle = Column(Enum(BillingCycle), nullable=False)  # "monthly" or "yearly"
    
    # Renewal Information
    renewal_type = Column(String(20), nullable=False)  # "new", "renewal", "upgrade", "downgrade"
    renewal_count = Column(Integer, default=0)  # Sequential renewal number (1st renewal, 2nd renewal, etc.)
    
    # Previous Subscription Snapshot (before renewal)
    previous_plan_name = Column(String(50), nullable=True)
    previous_plan_tier = Column(String(20), nullable=True)
    
    # Usage Summary Before Renewal (snapshot)
    usage_before_renewal = Column(JSON, nullable=True)  # Snapshot of usage before renewal
    
    # Payment Information
    payment_amount = Column(Float, default=0.0)
    payment_status = Column(String(20), default="pending")  # "pending", "paid", "failed"
    payment_date = Column(DateTime, nullable=True)
    stripe_invoice_id = Column(String(100), nullable=True)
    
    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow)
    
    # Relationships
    plan = relationship("SubscriptionPlan")
    
    # Indexes for performance
    __table_args__ = (
        {'mysql_engine': 'InnoDB'},
    )