workflow_profiler.py

#!/usr/bin/env python3
"""
Workflow Profiler CLI - AI-powered workflow resource profiling.

Analyzes HyperFlow workflow DAGs and recommends Kubernetes resource allocations
using an LLM-powered research and reasoning loop.
"""

import sys
from pathlib import Path

import click

from workflow_profiler.agent import WorkflowProfilerAgent
from workflow_profiler.output import generate_yaml_output, generate_markdown_report, generate_outputs
from workflow_profiler.config import GOOGLE_API_KEY, TAVILY_API_KEY
from workflow_profiler.tracer import enable_tracing, disable_tracing


def print_banner():
    """Print the Workflow Profiler banner."""
    click.echo(click.style("""
╔═══════════════════════════════════════════════════════════════╗
║                   Workflow Profiler Agent                     ║
║          AI-Powered Workflow Resource Profiling               ║
╚═══════════════════════════════════════════════════════════════╝
""", fg="cyan"))


def check_api_keys():
    """Check if required API keys are set."""
    if not GOOGLE_API_KEY:
        click.echo(click.style(
            "Error: GOOGLE_API_KEY not set. Please set it in .env file or environment.",
            fg="red"
        ))
        return False

    if not TAVILY_API_KEY:
        click.echo(click.style(
            "Warning: TAVILY_API_KEY not set. Web search will be disabled.",
            fg="yellow"
        ))

    return True


@click.group()
@click.version_option(version="0.1.0")
def cli():
    """Workflow Profiler - AI-powered workflow resource profiling."""
    pass


@cli.command()
@click.argument("workflow_path", type=click.Path(exists=True))
@click.option(
    "--output-dir", "-o",
    type=click.Path(),
    default=None,
    help="Output directory for reports (default: workflow directory)"
)
@click.option(
    "--format", "-f",
    type=click.Choice(["all", "yaml", "markdown"]),
    default="all",
    help="Output format (default: all)"
)
@click.option(
    "--quiet", "-q",
    is_flag=True,
    help="Suppress progress output"
)
@click.option(
    "--verbose", "-v",
    is_flag=True,
    help="Print verbose tracing to stdout (logs are always saved to file)"
)
@click.option(
    "--no-kb",
    is_flag=True,
    help="Ignore knowledge base, force research from scratch"
)
def analyze(workflow_path: str, output_dir: str, format: str, quiet: bool, verbose: bool, no_kb: bool):
    """
    Analyze a HyperFlow workflow and generate resource recommendations.

    WORKFLOW_PATH: Path to the workflow.json file
    """
    if not quiet:
        print_banner()

    if not check_api_keys():
        sys.exit(1)

    workflow_path = Path(workflow_path)

    # Determine output directory
    if output_dir:
        output_dir = Path(output_dir)
    else:
        output_dir = workflow_path.parent

    # Always enable tracing to log file, optionally print to stdout
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = output_dir / f"workflow_profiler_trace_{timestamp}.log"
    enable_tracing(log_file, print_to_stdout=verbose)

    if not quiet:
        click.echo(f"Analyzing workflow: {click.style(str(workflow_path), fg='green')}")
        click.echo(f"Output directory: {click.style(str(output_dir), fg='green')}")
        click.echo(f"Trace log: {click.style(str(log_file), fg='green')}")
        if no_kb:
            click.echo(click.style("Knowledge base disabled - researching from scratch", fg='yellow'))
        click.echo()

    # Progress callback
    def on_phase(phase: str, state: dict):
        if not quiet:
            phase_names = {
                "ingest": "Parsing workflow and discovering materials",
                "describe": "Understanding workflow tasks and structure",
                "analyze": "Analyzing workflow structure",
                "plan": "Planning research strategy",
                "gather": "Gathering evidence from sources",
                "synthesize": "Synthesizing task profiles",
                "verify": "Verifying claims and checking plausibility",
                "decide": "Making final recommendations",
                "complete": "Analysis complete",
            }
            phase_display = phase_names.get(phase, phase)
            click.echo(f"  → {click.style(phase_display, fg='cyan')}")

    # Run the agent
    try:
        agent = WorkflowProfilerAgent(use_knowledge_base=not no_kb)

        if not quiet:
            click.echo(click.style("\nRunning analysis...\n", bold=True))

        report = agent.run_with_callbacks(
            workflow_path,
            on_phase=on_phase if not quiet else None,
        )

        if not quiet:
            click.echo()

        # Generate outputs
        base_name = workflow_path.stem + "_resources"

        if format == "all":
            yaml_path, md_path = generate_outputs(report, output_dir, base_name)
            if not quiet:
                click.echo(click.style("\nOutputs generated:", bold=True))
                click.echo(f"  YAML:     {click.style(str(yaml_path), fg='green')}")
                click.echo(f"  Markdown: {click.style(str(md_path), fg='green')}")

        elif format == "yaml":
            yaml_path = output_dir / f"{base_name}.yaml"
            generate_yaml_output(report, yaml_path)
            if not quiet:
                click.echo(f"\nYAML output: {click.style(str(yaml_path), fg='green')}")

        elif format == "markdown":
            md_path = output_dir / f"{base_name}.md"
            generate_markdown_report(report, md_path)
            if not quiet:
                click.echo(f"\nMarkdown output: {click.style(str(md_path), fg='green')}")

        # Print summary
        if not quiet:
            click.echo()
            click.echo(click.style("═" * 60, fg="cyan"))
            click.echo(click.style("Analysis Summary", bold=True))
            click.echo(click.style("═" * 60, fg="cyan"))
            click.echo(f"  Workflow:         {report.get('workflow_name', 'unknown')}")
            click.echo(f"  Domain:           {report.get('domain', 'unknown')}")
            click.echo(f"  Task types:       {report.get('unique_task_types', 0)}")
            click.echo(f"  Total tasks:      {report.get('total_tasks', 0)}")
            click.echo(f"  Max parallelism:  {report.get('max_parallelism', 0)}")
            click.echo(f"  Recommendations:  {len(report.get('recommendations', {}))}")
            click.echo()

    except FileNotFoundError as e:
        click.echo(click.style(f"Error: {e}", fg="red"))
        sys.exit(1)
    except Exception as e:
        click.echo(click.style(f"Error during analysis: {e}", fg="red"))
        if not quiet:
            import traceback
            click.echo(traceback.format_exc())
        sys.exit(1)


@cli.command()
@click.argument("workflow_path", type=click.Path(exists=True))
def inspect(workflow_path: str):
    """
    Inspect a workflow without running full analysis.

    Shows workflow structure, task types, and available materials.
    """
    print_banner()

    from workflow_profiler.parser import parse_workflow, extract_task_types, compute_graph_metrics, discover_materials
    from workflow_profiler.knowledge_base import get_knowledge_base

    workflow_path = Path(workflow_path)

    click.echo(f"Inspecting: {click.style(str(workflow_path), fg='green')}")
    click.echo()

    # Parse workflow
    graph, metadata = parse_workflow(workflow_path)
    task_types = extract_task_types(graph)
    metrics = compute_graph_metrics(graph)
    materials = discover_materials(workflow_path.parent)

    # Display info
    click.echo(click.style("Workflow Metadata", bold=True))
    click.echo(f"  Name: {metadata.get('name', 'unknown')}")
    click.echo()

    click.echo(click.style("Structural Metrics", bold=True))
    click.echo(f"  Total tasks: {metrics.total_tasks}")
    click.echo(f"  Unique task types: {metrics.unique_task_types}")
    click.echo(f"  Max parallelism: {metrics.max_parallelism}")
    click.echo(f"  Critical path length: {metrics.critical_path_length}")
    click.echo()

    click.echo(click.style("Task Types", bold=True))
    kb = get_knowledge_base()
    for task_type, processes in sorted(task_types.items()):
        kb_match = kb.lookup(task_type)
        if kb_match:
            status = click.style(f"[KB: {kb_match.name}]", fg="green")
        else:
            status = click.style("[Unknown]", fg="yellow")
        click.echo(f"  {task_type}: {len(processes)} instances {status}")
    click.echo()

    click.echo(click.style("Available Materials", bold=True))
    if materials:
        for mat in materials:
            click.echo(f"  [{mat['material_type']}] {mat['filename']}")
    else:
        click.echo("  No materials found")
    click.echo()


@cli.command()
def list_tools():
    """List known tools in the knowledge base."""
    print_banner()

    from workflow_profiler.knowledge_base import get_knowledge_base

    kb = get_knowledge_base()

    click.echo(click.style("Known Tools in Knowledge Base", bold=True))
    click.echo()

    # Group by domain
    domains = kb.list_domains()
    for domain in sorted(domains):
        click.echo(click.style(f"  {domain.upper()}", fg="cyan", bold=True))
        tools = kb.get_tools_by_domain(domain)
        for tool in sorted(tools, key=lambda t: t.name):
            conf = click.style(f"[{tool.confidence:.0%}]", fg="green" if tool.confidence >= 0.8 else "yellow")
            click.echo(f"    {tool.name}: {tool.category} {conf}")
            if tool.aliases:
                click.echo(f"      aliases: {', '.join(tool.aliases[:3])}")
        click.echo()


@cli.command()
@click.argument("tool_name")
def tool_info(tool_name: str):
    """Show detailed information about a tool."""
    from workflow_profiler.knowledge_base import get_knowledge_base

    kb = get_knowledge_base()
    profile = kb.lookup(tool_name)

    if not profile:
        click.echo(click.style(f"Tool '{tool_name}' not found in knowledge base.", fg="red"))
        similar = kb.get_similar_tools(tool_name, limit=3)
        if similar:
            click.echo("\nDid you mean:")
            for p in similar:
                click.echo(f"  - {p.name}")
        sys.exit(1)

    click.echo(click.style(f"\nTool: {profile.name}", bold=True))
    click.echo(f"  Domain: {profile.domain}")
    click.echo(f"  Category: {profile.category}")
    click.echo(f"  Description: {profile.description}")
    click.echo()

    if profile.aliases:
        click.echo(f"  Aliases: {', '.join(profile.aliases)}")

    click.echo()
    click.echo(click.style("Resource Profile:", bold=True))
    if profile.cpu_util is not None:
        click.echo(f"  CPU Utilization: {profile.cpu_util:.0%}")
    if profile.memory_mb is not None:
        click.echo(f"  Memory: {profile.memory_mb} MB")
    if profile.io_pattern:
        click.echo(f"  I/O Pattern: {profile.io_pattern}")
    click.echo(f"  Confidence: {profile.confidence:.0%}")

    if profile.sources:
        click.echo()
        click.echo(click.style("Sources:", bold=True))
        for source in profile.sources:
            click.echo(f"  - {source}")


def main():
    """Main entry point."""
    cli()


if __name__ == "__main__":
    main()