"""Organization tools: tms_list_organizations."""

import asyncio
import json

from tms_mcp.api_client import ApiError, client


def register(mcp):
    @mcp.tool()
    async def tms_list_organizations(include_stats: bool = False) -> str:
        """List all organizations the current user belongs to.

        Args:
            include_stats: If True, includes task statistics (total, done, in_progress, blocked) for each org.
        """
        try:
            orgs = await client.get("/organizations/")
            if not orgs:
                return "No organizations found."

            if include_stats:
                stats_tasks = [
                    client.get(f"/organizations/{org['id']}/stats") for org in orgs
                ]
                stats_results = await asyncio.gather(*stats_tasks, return_exceptions=True)
                for org, stats in zip(orgs, stats_results):
                    if isinstance(stats, dict):
                        org["stats"] = stats

            return json.dumps(orgs, indent=2, default=str)
        except ApiError as e:
            return str(e)
