nowledge-graph-server

"""License management API endpoints."""

from typing import Optional
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
import real_ladybug as kuzu
import structlog

from ..dependencies import get_kuzu_client, KuzuClient
from ...services.license import get_license_service

logger = structlog.get_logger(__name__)

router = APIRouter(prefix="/api/license", tags=["license"])


class LicenseStatus(BaseModel):
    """License status information."""

    tier: str  # "free" or "pro"
    email: Optional[str] = None
    memory_limit: int
    memory_count: int
    features: dict
    # New fields for better state management
    is_device_activated: bool = False
    activation_status: str = "not_activated"  # "activated", "expired", "not_activated"
    device_id: Optional[str] = None


class LicenseActivateRequest(BaseModel):
    """Request to activate a license."""
    license_code: str  # base64-encoded license JSON
    email: str
    device_name: Optional[str] = None


class LicenseActivateResponse(BaseModel):
    """Response from license activation."""
    status: str  # "activated", "verified_offline", "error"
    message: str


class LicenseDeactivateResponse(BaseModel):
    """Response from license deactivation."""
    success: bool
    removed: list[str] = []


class LicenseRenewResponse(BaseModel):
    """Response from license renewal."""
    success: bool
    message: str


class LicenseError(HTTPException):
    """License-related error."""

    def __init__(self, detail: str, status_code: int = 402):
        super().__init__(status_code=status_code, detail=detail)


@router.get("/status", response_model=LicenseStatus, include_in_schema=False)
async def get_license_status(
    kuzu_client: KuzuClient = Depends(get_kuzu_client),
) -> LicenseStatus:
    """Get current license status and limits."""
    # Get memory count directly from database
    memory_count = get_memory_count(kuzu_client)

    # Get license status from the license file (written by Rust)
    license_service = get_license_service()
    tier = license_service.get_tier()
    email = license_service.get_email()
    is_device_activated = license_service.is_device_activated()
    activation_status = license_service.get_activation_status()

    # Define tier limits based on tier AND activation status
    # Pro features require both valid license AND device activation
    has_full_pro_access = tier == "pro" and is_device_activated

    if not has_full_pro_access:
        memory_limit = 20
        features = {
            "max_memories": 20,
            "remote_ai_models": False,
            "advanced_search": True,
            "knowledge_graph": True,
            "memory_distillation": True,
            "thread_import": True,
        }
    else:  # pro with activated device
        memory_limit = -1  # unlimited
        features = {
            "max_memories": -1,
            "remote_ai_models": True,
            "advanced_search": True,
            "knowledge_graph": True,
            "memory_distillation": True,
            "thread_import": True,
        }

    return LicenseStatus(
        tier=tier,
        email=email,
        memory_limit=memory_limit,
        memory_count=memory_count,
        features=features,
        is_device_activated=is_device_activated,
        activation_status=activation_status,
        device_id=license_service.get_device_id(),
    )


def get_memory_count(kuzu_client: KuzuClient) -> int:
    """Get current memory count from database.

    Args:
        kuzu_client: Kuzu database client

    Returns:
        Current number of memories in the database
    """
    query = "MATCH (m:Memory) RETURN COUNT(m) as count"
    result = kuzu_client.conn.execute(query)  # type: ignore[union-attr]
    assert isinstance(result, kuzu.QueryResult)
    return result.get_next()[0] if result.has_next() else 0


def check_memory_creation_limit(
    kuzu_client: KuzuClient, memories_to_create: int = 1
) -> None:
    """Check if creating new memories would exceed license limits.

    This function combines memory count retrieval, tier checking, and limit validation.

    Args:
        kuzu_client: Kuzu database client
        memories_to_create: Number of memories about to be created (default: 1)

    Raises:
        LicenseError: If creating the memories would exceed the tier limit
    """
    # Get current memory count from database
    current_count = get_memory_count(kuzu_client)

    # Get tier from license service
    from ...services.license import get_license_service

    tier = get_license_service().get_tier()

    # Check if creating these memories would exceed the limit
    check_memory_limit(current_count + memories_to_create, tier)


def check_memory_limit(current_count: int, tier: str = "free") -> None:
    """Check if memory limit has been reached.

    Args:
        current_count: Current number of memories
        tier: License tier ("free" or "pro")

    Raises:
        LicenseError: If memory limit is exceeded
    """
    if tier == "free" and current_count > 20:
        raise LicenseError(
            detail=f"Memory limit reached. Free tier allows up to 20 memories (current: {current_count}). Upgrade to Pro for unlimited memories.",
            status_code=402,  # Payment Required
        )


# =========================================================================
# License activation / deactivation / renewal (headless / server support)
# =========================================================================


@router.post("/activate", response_model=LicenseActivateResponse, include_in_schema=False)
async def activate_license(request: LicenseActivateRequest) -> LicenseActivateResponse:
    """Activate a license key.

    Verifies the license signature, checks email match, calls the remote
    license server for device activation, and saves license + device auth
    files to disk.
    """
    license_service = get_license_service()
    result = await license_service.activate_license(
        license_code_b64=request.license_code,
        email=request.email,
        device_name=request.device_name,
    )
    logger.info("License activation result", status=result["status"])
    return LicenseActivateResponse(**result)


@router.post("/deactivate", response_model=LicenseDeactivateResponse, include_in_schema=False)
async def deactivate_license() -> LicenseDeactivateResponse:
    """Deactivate the current license (remove license and device auth files)."""
    license_service = get_license_service()
    result = license_service.deactivate_license()
    logger.info("License deactivated", removed=result["removed"])
    return LicenseDeactivateResponse(**result)


@router.post("/renew", response_model=LicenseRenewResponse, include_in_schema=False)
async def renew_license() -> LicenseRenewResponse:
    """Renew device authorization (30-day rolling auth check-in)."""
    license_service = get_license_service()
    result = await license_service.renew_activation()
    logger.info("License renewal result", success=result["success"])
    return LicenseRenewResponse(**result)