nexus-5/core/mcp/tools/dashboard.py

"""Dashboard tools for MCP."""

from datetime import date, timedelta
from typing import Optional

from channels.db import database_sync_to_async

from core.mcp.auth import MCPContext, Role
from core.mcp.base import mcp, json_response, error_response, logger


@mcp.tool()
async def get_my_schedule(
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    status: Optional[str] = None
) -> str:
    """
    Get your assigned services and projects for a date range.

    Args:
        start_date: Start date in YYYY-MM-DD format (defaults to today)
        end_date: End date in YYYY-MM-DD format (defaults to 7 days from start)
        status: Optional status filter (SCHEDULED, IN_PROGRESS, COMPLETED)

    Returns:
        JSON object with services and projects arrays
    """
    profile = MCPContext.get_profile()
    if not profile:
        return error_response("No active profile. Call set_active_profile first.")

    from datetime import datetime

    # Default date range
    if not start_date:
        start_date = date.today().isoformat()
    if not end_date:
        start = datetime.strptime(start_date, "%Y-%m-%d").date()
        end_date = (start + timedelta(days=7)).isoformat()

    @database_sync_to_async
    def fetch_schedule():
        from core.models import Service, Project

        # Build base querysets
        if profile.role == Role.TEAM_MEMBER.value:
            services_qs = Service.objects.filter(team_members__id=profile.id)
            projects_qs = Project.objects.filter(team_members__id=profile.id)
        else:
            services_qs = Service.objects.all()
            projects_qs = Project.objects.all()

        # Apply date filters
        services_qs = services_qs.filter(date__gte=start_date, date__lte=end_date)
        projects_qs = projects_qs.filter(date__gte=start_date, date__lte=end_date)

        # Apply status filter
        if status:
            services_qs = services_qs.filter(status=status)
            projects_qs = projects_qs.filter(status=status)

        # Fetch with related data
        services_qs = services_qs.select_related(
            'account_address__account__customer'
        ).prefetch_related('team_members').order_by('date')

        projects_qs = projects_qs.select_related(
            'customer', 'account_address__account'
        ).prefetch_related('team_members').order_by('date')

        services = []
        for s in services_qs[:50]:
            addr = s.account_address
            services.append({
                "id": str(s.id),
                "type": "service",
                "date": str(s.date),
                "status": s.status,
                "customer": addr.account.customer.name,
                "account": addr.account.name,
                "location": addr.name or "Primary",
                "address": f"{addr.street_address}, {addr.city}, {addr.state} {addr.zip_code}",
                "team_members": [
                    f"{t.first_name} {t.last_name}".strip()
                    for t in s.team_members.all() if t.role != 'ADMIN'
                ]
            })

        projects = []
        for p in projects_qs[:50]:
            if p.account_address:
                addr = p.account_address
                location = addr.name or "Primary"
                address = f"{addr.street_address}, {addr.city}, {addr.state} {addr.zip_code}"
                account = addr.account.name
            else:
                location = None
                address = f"{p.street_address}, {p.city}, {p.state} {p.zip_code}"
                account = None

            projects.append({
                "id": str(p.id),
                "type": "project",
                "name": p.name,
                "date": str(p.date),
                "status": p.status,
                "customer": p.customer.name,
                "account": account,
                "location": location,
                "address": address,
                "labor": float(p.labor),
                "amount": float(p.amount),
                "team_members": [
                    f"{t.first_name} {t.last_name}".strip()
                    for t in p.team_members.all() if t.role != 'ADMIN'
                ]
            })

        return {"services": services, "projects": projects}

    try:
        result = await fetch_schedule()
        return json_response(result)
    except Exception as e:
        logger.error(f"Error fetching schedule: {e}")
        return error_response(str(e))


@mcp.tool()
async def get_system_stats() -> str:
    """
    Get high-level system statistics. Requires ADMIN or TEAM_LEADER role.

    Returns:
        JSON object with counts of customers, accounts, services, projects, etc.
    """
    profile = MCPContext.get_profile()
    if not profile:
        return error_response("No active profile. Call set_active_profile first.")

    if profile.role == Role.TEAM_MEMBER.value:
        return error_response("Access denied. TEAM_LEADER or ADMIN role required.")

    from core.models import Customer, Account, Service, Project, TeamProfile

    @database_sync_to_async
    def fetch_stats():
        return {
            "customers": Customer.objects.count(),
            "accounts": Account.objects.count(),
            "services": {
                "total": Service.objects.count(),
                "scheduled": Service.objects.filter(status='SCHEDULED').count(),
                "in_progress": Service.objects.filter(status='IN_PROGRESS').count(),
                "completed": Service.objects.filter(status='COMPLETED').count(),
            },
            "projects": {
                "total": Project.objects.count(),
                "scheduled": Project.objects.filter(status='SCHEDULED').count(),
                "in_progress": Project.objects.filter(status='IN_PROGRESS').count(),
                "completed": Project.objects.filter(status='COMPLETED').count(),
            },
            "team_members": TeamProfile.objects.exclude(role='ADMIN').count(),
        }

    stats = await fetch_stats()
    return json_response(stats)