tf_code/packages/tfcode/python/tf_sync/config.py

"""
Configuration management for tfcode ToothFairyAI integration.
Uses the official ToothFairyAI Python SDK for multi-region support.
"""

import os
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field, SecretStr
from toothfairyai import ToothFairyClient
from toothfairyai.errors import ToothFairyError


class Region(str, Enum):
    DEV = "dev"
    AU = "au"
    EU = "eu"
    US = "us"


class ToolType(str, Enum):
    MCP_SERVER = "mcp_server"
    AGENT_SKILL = "agent_skill"
    CODER_AGENT = "coder_agent"
    DATABASE_SCRIPT = "database_script"
    API_FUNCTION = "api_function"
    PROMPT = "prompt"


class FunctionRequestType(str, Enum):
    GET = "get"
    POST = "post"
    PUT = "put"
    DELETE = "delete"
    PATCH = "patch"
    CUSTOM = "custom"
    GRAPHQL_QUERY = "graphql_query"
    GRAPHQL_MUTATION = "graphql_mutation"


# Region-specific URL configurations
REGION_URLS = {
    Region.DEV: {
        "base_url": "https://api.toothfairylab.link",
        "ai_url": "https://ai.toothfairylab.link",
        "ai_stream_url": "https://ais.toothfairylab.link",
        "mcp_url": "https://mcp.toothfairylab.link/sse",
        "mcp_proxy_url": "https://mcp-proxy.toothfairylab.link",
    },
    Region.AU: {
        "base_url": "https://api.toothfairyai.com",
        "ai_url": "https://ai.toothfairyai.com",
        "ai_stream_url": "https://ais.toothfairyai.com",
        "mcp_url": "https://mcp.toothfairyai.com/sse",
        "mcp_proxy_url": "https://mcp-proxy.toothfairyai.com",
    },
    Region.EU: {
        "base_url": "https://api.eu.toothfairyai.com",
        "ai_url": "https://ai.eu.toothfairyai.com",
        "ai_stream_url": "https://ais.eu.toothfairyai.com",
        "mcp_url": "https://mcp.eu.toothfairyai.com/sse",
        "mcp_proxy_url": "https://mcp-proxy.eu.toothfairyai.com",
    },
    Region.US: {
        "base_url": "https://api.us.toothfairyai.com",
        "ai_url": "https://ai.us.toothfairyai.com",
        "ai_stream_url": "https://ais.us.toothfairyai.com",
        "mcp_url": "https://mcp.us.toothfairyai.com/sse",
        "mcp_proxy_url": "https://mcp-proxy.us.toothfairyai.com",
    },
}


def get_region_urls(region: Region) -> dict[str, str]:
    """Get URLs for a specific region."""
    return REGION_URLS.get(region, REGION_URLS[Region.AU])


class TFConfig(BaseModel):
    """ToothFairyAI workspace configuration."""
    
    workspace_id: str
    api_key: SecretStr
    region: Region = Region.AU
    enabled: bool = True
    
    sync_interval: int = Field(default=3600, ge=60)
    mcp_proxy_timeout: int = Field(default=30000, ge=1000)
    
    _client: Optional[ToothFairyClient] = None
    
    def get_client(self) -> ToothFairyClient:
        """
        Get or create a ToothFairyClient instance configured for this region.
        
        Returns:
            ToothFairyClient configured with region-specific URLs
        """
        if self._client is None:
            urls = get_region_urls(self.region)
            self._client = ToothFairyClient(
                api_key=self.api_key.get_secret_value(),
                workspace_id=self.workspace_id,
                base_url=urls["base_url"],
                ai_url=urls["ai_url"],
                ai_stream_url=urls["ai_stream_url"],
            )
        return self._client
    
    @property
    def mcp_sse_url(self) -> str:
        """Get the MCP SSE endpoint URL for this region."""
        return get_region_urls(self.region)["mcp_url"]
    
    @property
    def mcp_proxy_url(self) -> str:
        """Get the MCP proxy URL for this region."""
        return get_region_urls(self.region)["mcp_proxy_url"]


class CredentialValidationResult(BaseModel):
    """Result of credential validation."""
    success: bool
    workspace_id: Optional[str] = None
    workspace_name: Optional[str] = None
    error: Optional[str] = None


def load_config(
    workspace_id: Optional[str] = None,
    api_key: Optional[str] = None,
    region: Optional[Region] = None,
) -> TFConfig:
    """
    Load ToothFairyAI configuration from environment or parameters.
    
    Args:
        workspace_id: Workspace UUID (defaults to TF_WORKSPACE_ID env var)
        api_key: API key (defaults to TF_API_KEY env var)
        region: Region (defaults to TF_REGION env var or 'au')
    
    Returns:
        TFConfig instance
    
    Raises:
        ValueError: If required configuration is missing
    """
    ws_id = workspace_id or os.environ.get("TF_WORKSPACE_ID")
    key = api_key or os.environ.get("TF_API_KEY")
    
    # Parse region from env or use provided/default
    region_str = os.environ.get("TF_REGION", "au")
    reg = region or Region(region_str)
    
    if not ws_id:
        raise ValueError("TF_WORKSPACE_ID not set. Set environment variable or pass workspace_id.")
    if not key:
        raise ValueError("TF_API_KEY not set. Set environment variable or pass api_key.")
    
    return TFConfig(
        workspace_id=ws_id,
        api_key=SecretStr(key),
        region=reg,
    )


def validate_credentials(config: TFConfig) -> CredentialValidationResult:
    """
    Validate ToothFairyAI credentials using the SDK.
    
    Args:
        config: TFConfig instance
    
    Returns:
        CredentialValidationResult indicating success or failure
    """
    try:
        client = config.get_client()
        
        # Test connection by listing chats (lightweight operation)
        if client.test_connection():
            return CredentialValidationResult(
                success=True,
                workspace_id=config.workspace_id,
                workspace_name="Connected",
            )
        else:
            return CredentialValidationResult(
                success=False,
                error="Connection test failed. Check credentials and region.",
            )
            
    except ToothFairyError as e:
        error_msg = str(e)
        
        if "401" in error_msg or "Unauthorized" in error_msg:
            return CredentialValidationResult(
                success=False,
                error="Invalid API key. Check TF_API_KEY environment variable.",
            )
        elif "403" in error_msg or "Forbidden" in error_msg:
            return CredentialValidationResult(
                success=False,
                error="API access not allowed. Business or Enterprise subscription required.",
            )
        elif "404" in error_msg or "Not Found" in error_msg:
            return CredentialValidationResult(
                success=False,
                error="Workspace not found. Check TF_WORKSPACE_ID environment variable.",
            )
        else:
            return CredentialValidationResult(
                success=False,
                error=f"API error: {error_msg}",
            )
    except Exception as e:
        return CredentialValidationResult(
            success=False,
            error=f"Unexpected error: {str(e)}",
        )