import fire
import ast
import importlib
from pathlib import Path
from rich.console import Console


HERE = Path(__file__).parent

console = Console()


def _toolset_class_name(module_name: str) -> str:
    acronyms = {"api": "API", "rag": "RAG", "scfm": "SCFM"}
    return "".join(
        acronyms.get(word.lower(), word.capitalize())
        for word in module_name.split("_")
    ) + "ToolSet"


def get_toolset_modules():
    return [
        f.stem for f in HERE.glob("*")
        if (not f.stem.startswith("__")) and (f.stem) not in ["bio"]
    ]


def extract_exported_classes(content: str):
    """Extract exported classes from a module."""
    tree = ast.parse(content)
    classes = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign):
            if isinstance(node.targets[0], ast.Name) and node.targets[0].id == "__all__":
                classes.extend(node.value.elts)
    return [
        c.value if isinstance(c, ast.Constant) else c.s
        for c in classes
    ]


def import_toolset_class(module_name: str):
    module_path = HERE / module_name
    if module_path.is_dir():
        content = (module_path / "__init__.py").read_text(encoding="utf-8")
    else:
        content = (HERE / f"{module_name}.py").read_text(encoding="utf-8")
    classes = extract_exported_classes(content)
    toolset_classes = [class_name for class_name in classes if class_name.endswith("ToolSet")]
    if len(toolset_classes) == 1:
        class_name = toolset_classes[0]
    elif _toolset_class_name(module_name) in toolset_classes:
        class_name = _toolset_class_name(module_name)
    else:
        raise ImportError(
            f"Module {module_name} has incorrect number of ToolSet classes: {toolset_classes}"
        )
    if class_name:
        mod = importlib.import_module(f"pantheon.toolsets.{module_name}")
        return getattr(mod, class_name)


def list_toolsets():
    """List all available built-in toolsets."""
    print()
    console.print("[bold]Available built-in toolsets:[/bold]")
    print()
    for module_name in sorted(get_toolset_modules()):
        toolset_class = import_toolset_class(module_name)
        doc = toolset_class.__doc__
        if doc:
            lines = doc.split("\n")
            for line in lines:
                if line.strip() != "":
                    break
            console.print(f"- [bold]{module_name}[/bold]: [dim]{line.strip()}[/dim]")
        else:
            console.print(f"- [bold]{module_name}[/bold]")
    print()


def detail(toolset_name: str):
    """View detailed information about a toolset."""
    toolset_class = import_toolset_class(toolset_name)
    doc = toolset_class.__doc__
    print()
    console.print(f"[bold]{toolset_class.__name__} Doc-string:[/bold]")
    print()
    console.print(doc)
    print()
    console.print("Start this toolset with:")
    print()
    console.print(f"[bold]$[/bold] [dim]python -m pantheon.toolsets start {toolset_name}[/dim]")
    print()


async def start(
    toolset_name: str,
    service_name: str = None,
    mcp: bool = False,
    mcp_kwargs: dict = {},
    **kwargs,
    ):
    """Start a remote toolset.

    Args:
        toolset_name: The name of the toolset to run.
        mcp: Whether to run the toolset as an MCP server.
        mcp_kwargs: The keyword arguments for the MCP server.
    """
    if service_name is None:
        service_name = toolset_name
    toolset_class = import_toolset_class(toolset_name)
    toolset = toolset_class(service_name, **kwargs)
    if mcp:
        await toolset.run_as_mcp(**mcp_kwargs)
    else:
        await toolset.run()


fire.Fire({
    "start": start,
    "list": list_toolsets,
    "detail": detail,
})
