tools/platform_management_tool.py

"""
Platform Management Tools for Atom
Provides administrative tools for managing tenants, workspaces, members, and teams.
"""

import logging
import uuid
import json
from typing import Any, Dict, List, Optional
from datetime import datetime

from core.database import SessionLocal
from core.models import Tenant, Workspace, User, Team, team_members
from sqlalchemy import or_

logger = logging.getLogger(__name__)

async def get_platform_settings(context: Dict[str, Any] = None) -> Dict[str, Any]:
    """Retrieves platform-wide settings and configuration for the current tenant."""
    from core.database import SessionLocal
    from core.models import TenantSetting
    
    workspace_id = (context.get("workspace_id") if context else None) or "default"
        
    try:
        with SessionLocal() as db:
            settings = db.query(TenantSetting).filter(TenantSetting.tenant_id == workspace_id).all()
            return {s.setting_key: s.setting_value for s in settings}
    except Exception as e:
        logger.error(f"Error fetching platform settings: {e}")
        return {"error": f"Failed to fetch settings: {str(e)}"}

async def update_platform_setting(key: str, value: str, context: Dict[str, Any] = None) -> str:
    """Updates or creates a platform-wide setting."""
    from core.database import SessionLocal
    from core.models import TenantSetting
    
    workspace_id = (context.get("workspace_id") if context else None) or "default"
        
    try:
        with SessionLocal() as db:
            setting = db.query(TenantSetting).filter(
                TenantSetting.tenant_id == workspace_id,
                TenantSetting.setting_key == key
            ).first()
            
            if setting:
                setting.setting_value = value
            else:
                setting = TenantSetting(tenant_id=workspace_id, setting_key=key, setting_value=value)
                db.add(setting)
            
            db.commit()
            return f"Setting '{key}' successfully updated to '{value}'."
    except Exception as e:
        logger.error(f"Error updating platform setting: {e}")
        return f"Error: Failed to update setting: {str(e)}"

async def update_tenant_profile(
    name: Optional[str] = None,
    billing_email: Optional[str] = None,
    logo_url: Optional[str] = None,
    primary_color: Optional[str] = None,
    budget_limit_usd: Optional[float] = None,
    context: Dict[str, Any] = None
) -> str:
    """Updates the tenant's profile information."""
    from core.database import SessionLocal
    from core.models import Tenant, Workspace
    
    workspace_id = (context.get("workspace_id") if context else None) or "default"
        
    try:
        with SessionLocal() as db:
            # First get tenant_id from workspace
            ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
            
            # Fallback for single-tenant mode where workspace might not exist yet
            tenant_id = ws.tenant_id if ws else "default"
            
            tenant = db.query(Tenant).filter(Tenant.id == tenant_id).first()
            
            # If default tenant doesn't exist, we might be in ultra-lean OS mode
            if not tenant and tenant_id == "default":
                # Create default tenant on the fly if needed? Or just return error
                return f"Error: Default tenant not found. Please initialize the system."
            elif not tenant:
                return f"Error: Tenant {tenant_id} not found."
            
            updates = []
            if name:
                tenant.name = name
                updates.append("name")
            if billing_email:
                tenant.billing_email = billing_email
                updates.append("billing_email")
            if logo_url:
                metadata = tenant.metadata_json or {}
                metadata["logo_url"] = logo_url
                tenant.metadata_json = metadata
                updates.append("logo_url")
            if primary_color:
                metadata = tenant.metadata_json or {}
                metadata["primary_color"] = primary_color
                tenant.metadata_json = metadata
                updates.append("primary_color")
            if budget_limit_usd is not None:
                tenant.budget_limit_usd = budget_limit_usd
                updates.append("budget_limit_usd")
                
            if not updates:
                return "No updates provided."
                
            db.commit()
            return f"Tenant profile updated successfully. Fields modified: {', '.join(updates)}."
    except Exception as e:
        logger.error(f"Error updating tenant profile: {e}")
        return f"Error: Failed to update tenant profile: {str(e)}"

async def set_byok_api_key(
    provider: str,
    api_key: str,
    context: Dict[str, Any] = None
) -> str:
    """Sets or updates a Bring-Your-Own-Key (BYOK) API key for a provider."""
    from core.database import SessionLocal
    from core.byok_endpoints import BYOKManager
    
    tenant_id = context.get("workspace_id") if context else None
    if not tenant_id:
        return "Error: Could not resolve tenant/workspace ID from context."
        
    db = SessionLocal()
    try:
        manager = BYOKManager(db)
        await manager.set_api_key(tenant_id, provider, api_key)
        db.commit()
        return f"Successfully set API key for {provider} in tenant {tenant_id}."
    except Exception as e:
        db.rollback()
        logger.error(f"Error setting BYOK API key: {e}")
        return f"Error setting BYOK API key: {str(e)}"
    finally:
        db.close()

async def list_tenant_members(context: Dict[str, Any] = None) -> str:
    """Lists all members of the current tenant."""
    from core.database import SessionLocal
    from core.models import User, Workspace
    
    workspace_id = context.get("workspace_id") if context else None
    if not workspace_id:
        return "Error: Could not resolve workspace ID from context."
        
    db = SessionLocal()
    try:
        # First get the tenant_id from the workspace
        ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
        if not ws:
            return f"Error: Workspace {workspace_id} not found."
            
        tenant_id = ws.tenant_id
        
        # Get all users belonging to this tenant
        # In multi-tenant architecture, users usually have a tenant_id
        members = db.query(User).filter(User.tenant_id == tenant_id).all()
        
        if not members:
            return "No members found for this tenant."
            
        result = [f"Members for Tenant {tenant_id}:"]
        for m in members:
            result.append(f"- {m.full_name or m.email} (ID: {m.id}, Role: {getattr(m, 'role', 'N/A')}, Status: {getattr(m, 'status', 'N/A')})")
            
        return "\n".join(result)
    except Exception as e:
        logger.error(f"Error listing tenant members: {e}")
        return f"Error listing tenant members: {str(e)}"
    finally:
        db.close()

async def manage_tenant_member(
    user_id: str,
    action: str, # "update_role", "deactivate", "reactivate"
    role: Optional[str] = None,
    context: Dict[str, Any] = None
) -> str:
    """Updates a member's role or status within a tenant."""
    from core.database import SessionLocal
    from core.models import User
    
    db = SessionLocal()
    try:
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return f"Error: User {user_id} not found."
            
        if action == "update_role":
            if not role:
                return "Error: role is required for update_role action."
            user.role = role
            msg = f"User {user_id} role updated to {role}."
        elif action == "deactivate":
            user.is_active = False
            msg = f"User {user_id} deactivated."
        elif action == "reactivate":
            user.is_active = True
            msg = f"User {user_id} reactivated."
        else:
            return f"Error: Unknown action {action}."
            
        db.commit()
        return msg
    except Exception as e:
        db.rollback()
        logger.error(f"Error managing tenant member: {e}")
        return f"Error managing tenant member: {str(e)}"
    finally:
        db.close()

async def manage_workspace(
    name: str,
    action: str = "create", # "create", "update"
    workspace_id: Optional[str] = None,
    description: Optional[str] = None,
    is_startup: bool = False,
    context: Dict[str, Any] = None
) -> str:
    """Creates or updates a workspace within a tenant."""
    from core.database import SessionLocal
    from core.models import Workspace, Tenant
    
    # Resolve tenant_id
    actual_tenant_id = context.get("tenant_id") if context else None
    if not actual_tenant_id:
        workspace_id_ctx = context.get("workspace_id") if context else None
        if workspace_id_ctx:
            # Try to lookup tenant from workspace
            with SessionLocal() as db_temp:
                ws_temp = db_temp.query(Workspace).filter(Workspace.id == workspace_id_ctx).first()
                if ws_temp:
                    actual_tenant_id = ws_temp.tenant_id
        
    if not actual_tenant_id:
        return "Error: Could not resolve tenant ID from context."

    db = SessionLocal()
    try:
        if action == "create":
            new_ws = Workspace(
                name=name,
                description=description,
                is_startup=is_startup,
                tenant_id=actual_tenant_id
            )
            db.add(new_ws)
            db.commit()
            return f"Workspace '{name}' created successfully with ID {new_ws.id}."
        
        elif action == "update":
            if not workspace_id:
                return "Error: workspace_id is required for update action."
            ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
            if not ws:
                return f"Error: Workspace {workspace_id} not found."
            
            ws.name = name
            if description is not None:
                ws.description = description
            ws.is_startup = is_startup
            db.commit()
            return f"Workspace {workspace_id} updated successfully."
        else:
            return f"Error: Unknown action {action}."
            
    except Exception as e:
        db.rollback()
        logger.error(f"Error managing workspace: {e}")
        return f"Error managing workspace: {str(e)}"
    finally:
        db.close()

async def manage_team(
    name: str,
    action: str = "create", # "create", "update"
    team_id: Optional[str] = None,
    workspace_id: Optional[str] = None,
    add_members: Optional[List[str]] = None, # List of user emails or IDs
    context: Dict[str, Any] = None
) -> str:
    """Creates or updates a Team and its members."""
    from core.database import SessionLocal
    from core.models import Team, User, team_members, Workspace, Tenant
    
    tenant_id = context.get("workspace_id") if context else None
    if not tenant_id:
        return "Error: Could not resolve tenant/workspace context."

    db = SessionLocal()
    try:
        # Resolve actual workspace_id if not provided
        if not workspace_id:
            workspace_id = context.get("workspace_id")

        if action == "create":
            team = Team(
                name=name,
                workspace_id=workspace_id
            )
            db.add(team)
            db.flush() # Get team.id
            team_id = team.id
            msg = f"Team '{name}' created successfully with ID {team_id}."
        elif action == "update":
            if not team_id:
                return "Error: team_id is required for update action."
            team = db.query(Team).filter(Team.id == team_id).first()
            if not team:
                return f"Error: Team {team_id} not found."
            team.name = name
            msg = f"Team {team_id} updated successfully."
        else:
            return f"Error: Unknown action {action}."

        # Add members if provided
        if add_members and team_id:
            added = 0
            for identifier in add_members:
                user = db.query(User).filter(or_(User.id == identifier, User.email == identifier)).first()
                if user:
                    # Check if already in team
                    existing = db.query(team_members).filter(
                        team_members.c.user_id == user.id,
                        team_members.c.team_id == team_id
                    ).first()
                    if not existing:
                        db.execute(team_members.insert().values(user_id=user.id, team_id=team_id))
                        added += 1
            msg += f" Added {added} members."

        db.commit()
        return msg
    except Exception as e:
        db.rollback()
        logger.error(f"Error managing team: {e}")
        return f"Error managing team: {str(e)}"
    finally:
        db.close()