app/core/api_gateway.py

"""
API Gateway Models - Groq/OpenRouter-style account and chat management
"""
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import List, Optional, Dict
from enum import Enum
import secrets
import hashlib


class AccountStatus(str, Enum):
    ACTIVE = "active"
    SUSPENDED = "suspended"
    TRIAL = "trial"


class Role(str, Enum):
    USER = "user"
    ADMIN = "admin"


@dataclass
class APIKey:
    key_id: str
    key_hash: str
    account_id: str
    name: str
    created_at: datetime = field(default_factory=datetime.utcnow)
    last_used: Optional[datetime] = None
    expires_at: Optional[datetime] = None
    is_active: bool = True
    rate_limit_per_minute: int = 60
    rate_limit_per_day: int = 1000
    
    @staticmethod
    def generate_key() -> str:
        """Generate a secure API key."""
        return f"m5_{secrets.token_urlsafe(32)}"
    
    @staticmethod
    def hash_key(key: str) -> str:
        """Hash an API key for storage."""
        return hashlib.sha256(key.encode()).hexdigest()


@dataclass
class ChatMessage:
    role: str  # "user", "assistant", "system"
    content: str
    timestamp: datetime = field(default_factory=datetime.utcnow)
    tokens_used: Optional[int] = None
    model: Optional[str] = None
    latency_ms: Optional[int] = None


@dataclass
class ChatSession:
    session_id: str
    account_id: str
    title: str
    messages: List[ChatMessage] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.utcnow)
    updated_at: datetime = field(default_factory=datetime.utcnow)
    model: Optional[str] = None
    total_tokens: int = 0
    total_cost_usd: float = 0.0
    
    def add_message(self, message: ChatMessage):
        self.messages.append(message)
        self.updated_at = datetime.utcnow()
        if message.tokens_used:
            self.total_tokens += message.tokens_used


@dataclass
class Account:
    account_id: str
    email: str
    password_hash: str
    name: str
    status: AccountStatus = AccountStatus.TRIAL
    role: Role = Role.USER
    created_at: datetime = field(default_factory=datetime.utcnow)
    api_keys: List[APIKey] = field(default_factory=list)
    chat_sessions: List[ChatSession] = field(default_factory=list)
    balance_usd: float = 10.0  # Free trial credits
    total_requests: int = 0
    total_tokens: int = 0
    
    @staticmethod
    def hash_password(password: str) -> str:
        """Hash a password for storage."""
        return hashlib.sha256(password.encode()).hexdigest()
    
    def verify_password(self, password: str) -> bool:
        """Verify a password against the hash."""
        return self.password_hash == self.hash_password(password)
    
    def add_api_key(self, name: str, expires_days: Optional[int] = None) -> APIKey:
        """Add a new API key to the account."""
        key = APIKey.generate_key()
        key_hash = APIKey.hash_key(key)
        expires_at = datetime.utcnow() + timedelta(days=expires_days) if expires_days else None
        
        api_key = APIKey(
            key_id=secrets.token_hex(8),
            key_hash=key_hash,
            account_id=self.account_id,
            name=name,
            expires_at=expires_at
        )
        self.api_keys.append(api_key)
        return api_key
    
    def revoke_api_key(self, key_id: str) -> bool:
        """Revoke an API key."""
        for key in self.api_keys:
            if key.key_id == key_id:
                key.is_active = False
                return True
        return False
    
    def create_chat_session(self, title: str) -> ChatSession:
        """Create a new chat session."""
        session = ChatSession(
            session_id=secrets.token_hex(16),
            account_id=self.account_id,
            title=title
        )
        self.chat_sessions.append(session)
        return session
    
    def deduct_balance(self, amount_usd: float) -> bool:
        """Deduct balance from account."""
        if self.balance_usd >= amount_usd:
            self.balance_usd -= amount_usd
            return True
        return False


@dataclass
class UsageMetrics:
    account_id: str
    date: datetime
    requests_count: int = 0
    tokens_count: int = 0
    cost_usd: float = 0.0
    model_usage: Dict[str, int] = field(default_factory=dict)
    
    def record_request(self, tokens: int, model: str, cost_usd: float):
        self.requests_count += 1
        self.tokens_count += tokens
        self.cost_usd += cost_usd
        self.model_usage[model] = self.model_usage.get(model, 0) + tokens


# In-memory storage (replace with database in production)
_accounts: Dict[str, Account] = {}
_api_keys: Dict[str, APIKey] = {}  # key_hash -> APIKey
_rate_limits: Dict[str, List[datetime]] = {}  # key_id -> request timestamps


class AccountManager:
    """Manage accounts and API keys."""
    
    @staticmethod
    def create_account(email: str, password: str, name: str) -> Account:
        """Create a new account."""
        account_id = secrets.token_hex(16)
        account = Account(
            account_id=account_id,
            email=email,
            password_hash=Account.hash_password(password),
            name=name
        )
        _accounts[account_id] = account
        return account
    
    @staticmethod
    def get_account_by_email(email: str) -> Optional[Account]:
        """Get account by email."""
        for account in _accounts.values():
            if account.email == email:
                return account
        return None
    
    @staticmethod
    def get_account_by_id(account_id: str) -> Optional[Account]:
        """Get account by ID."""
        return _accounts.get(account_id)
    
    @staticmethod
    def verify_api_key(api_key: str) -> Optional[Account]:
        """Verify an API key and return the associated account."""
        key_hash = APIKey.hash_key(api_key)
        for account in _accounts.values():
            for key in account.api_keys:
                if key.key_hash == key_hash and key.is_active:
                    if key.expires_at and datetime.utcnow() > key.expires_at:
                        key.is_active = False
                        return None
                    key.last_used = datetime.utcnow()
                    return account
        return None
    
    @staticmethod
    def check_rate_limit(api_key: str) -> bool:
        """Check if API key is within rate limits."""
        key_hash = APIKey.hash_key(api_key)
        for account in _accounts.values():
            for key in account.api_keys:
                if key.key_hash == key_hash:
                    now = datetime.utcnow()
                    minute_ago = now - timedelta(minutes=1)
                    day_ago = now - timedelta(days=1)
                    
                    if key.key_id not in _rate_limits:
                        _rate_limits[key.key_id] = []
                    
                    # Clean old timestamps
                    _rate_limits[key.key_id] = [
                        ts for ts in _rate_limits[key.key_id]
                        if ts > day_ago
                    ]
                    
                    # Check limits
                    recent_minute = sum(1 for ts in _rate_limits[key.key_id] if ts > minute_ago)
                    recent_day = len(_rate_limits[key.key_id])
                    
                    if recent_minute >= key.rate_limit_per_minute:
                        return False
                    if recent_day >= key.rate_limit_per_day:
                        return False
                    
                    # Record this request
                    _rate_limits[key.key_id].append(now)
                    return True
        return False
    
    @staticmethod
    def get_all_accounts() -> List[Account]:
        """Get all accounts (admin only)."""
        return list(_accounts.values())