#!/usr/bin/env python3
import json
import argparse
import os
from crewai import Agent, Task, Crew, Process
from langchain_openai import ChatOpenAI

def parse_arguments():
    parser = argparse.ArgumentParser(description='Generate a flow using CrewAI')
    parser.add_argument('--description', type=str, required=True, help='Description of the automation')
    parser.add_argument('--model', type=str, default='gpt-4', help='LLM model to use')
    parser.add_argument('--max_steps', type=int, default=10, help='Maximum number of steps')
    parser.add_argument('--output_path', type=str, required=True, help='Path to save the output JSON')
    return parser.parse_args()

def main():
    args = parse_arguments()
    
    # Check if OPENAI_API_KEY environment variable is set
    if 'OPENAI_API_KEY' not in os.environ:
        raise ValueError("OPENAI_API_KEY environment variable is required")

    # Initialize the LLM
    llm = ChatOpenAI(model=args.model)
    
    # Create the agents
    planner = Agent(
        role='Automation Planner',
        goal='Design efficient automation workflows based on user requirements',
        backstory='You are an expert in automation planning and breaking down complex tasks into logical steps.',
        llm=llm,
        verbose=True
    )
    
    developer = Agent(
        role='Automation Developer',
        goal='Convert high-level automation plans into detailed implementation steps',
        backstory='You specialize in translating conceptual plans into detailed technical implementations using the best tools and methods.',
        llm=llm,
        verbose=True
    )
    
    # Create the tasks
    planning_task = Task(
        description=f"Analyze this automation request and create a detailed plan: {args.description}. "
                    f"Break it down into clear, sequential steps, each with a specific goal.",
        agent=planner,
        expected_output="A numbered list of sequential steps to fulfill the automation request"
    )
    
    development_task = Task(
        description="For each step in the automation plan, provide detailed implementation instructions. "
                   "Each step should include: 1) The action type (browser, API, LLM, computation, or custom code), "
                   "2) Specific configuration needed, 3) Input and output details",
        agent=developer,
        expected_output="Detailed implementation steps with configurations for each step in the automation plan",
        context=[planning_task]
    )
    
    # Create and run the crew
    crew = Crew(
        agents=[planner, developer],
        tasks=[planning_task, development_task],
        verbose=2,
        process=Process.sequential
    )
    
    result = crew.kickoff()
    
    # Process the result
    steps = []
    
    # Parse the development task output to create steps
    lines = result.split('\n')
    current_step = None
    
    for line in lines:
        line = line.strip()
        if not line:
            continue
            
        # Look for step headers (Step 1:, etc.)
        if line.lower().startswith('step') and ':' in line:
            if current_step:
                steps.append(current_step)
            
            current_step = {
                'id': len(steps),
                'description': line.split(':', 1)[1].strip(),
                'action': '',
                'action_type': 'unknown',
                'config': {}
            }
        elif current_step and line.lower().startswith('action type:'):
            action_type = line.split(':', 1)[1].strip().lower()
            if 'browser' in action_type:
                current_step['action_type'] = 'browser'
            elif 'api' in action_type:
                current_step['action_type'] = 'api'
            elif 'llm' in action_type:
                current_step['action_type'] = 'llm'
            elif 'computation' in action_type:
                current_step['action_type'] = 'computation'
            elif 'custom' in action_type:
                current_step['action_type'] = 'custom'
            else:
                current_step['action_type'] = 'unknown'
        elif current_step and line.lower().startswith('config:'):
            current_step['config']['raw'] = line.split(':', 1)[1].strip()
        elif current_step and current_step.get('config'):
            # Add additional config details if it's a continuation
            if 'raw' in current_step['config']:
                current_step['config']['raw'] += '\n' + line
    
    # Add the last step
    if current_step:
        steps.append(current_step)
    
    # Generate the output
    output = {
        'description': args.description,
        'steps': steps,
        'raw_result': result
    }
    
    # Save to file
    with open(args.output_path, 'w') as f:
        json.dump(output, f, indent=2)
    
    print(f"Flow generation complete. Output saved to {args.output_path}")

if __name__ == "__main__":
    main()