import argparse
import subprocess
import json

def get_model_for_role(role: str) -> str:
    """Returns the specific model to use for a given tier role."""
    if role == 'tier1-orchestrator' or role == 'tier1':
        return 'gemini-3.1-pro-preview'
    elif role == 'tier2-tech-lead' or role == 'tier2':
        return 'gemini-3.0-flash-preview'
    else:
        return 'gemini-2.5-flash-lite'

def get_role_documents(role: str) -> list[str]:
    if role == 'tier1-orchestrator' or role == 'tier1':
        return ['conductor/product.md', 'conductor/product-guidelines.md']
    elif role == 'tier2-tech-lead' or role == 'tier2':
        return ['conductor/tech-stack.md', 'conductor/workflow.md']
    elif role == 'tier3-worker' or role == 'tier3':
        return ['conductor/workflow.md']
    return []

def execute_agent(role: str, prompt: str, docs: list[str]) -> str:
    model = get_model_for_role(role)
    command_text = f"Activate the mma-{role} skill. {prompt}"
    for doc in docs:
        command_text += f" @{doc}"
    
    cmd = ['gemini', '-p', command_text, '--output-format', 'json', '--model', model]
    try:
        process = subprocess.run(cmd, capture_output=True, text=True, shell=True)
        if not process.stdout and process.stderr:
            return f"Error: {process.stderr}"

        stdout = process.stdout
        start_index = stdout.find('{')
        if start_index != -1:
            json_str = stdout[start_index:]
            try:
                data = json.loads(json_str)
                return data.get('response', stdout)
            except json.JSONDecodeError:
                return stdout
        return stdout
    except Exception as e:
        return f"Execution failed: {str(e)}"

def create_parser():
    parser = argparse.ArgumentParser(description="MMA Execution Script")
    parser.add_argument(
        "--role", 
        choices=['tier1', 'tier2', 'tier3', 'tier4', 'tier1-orchestrator', 'tier2-tech-lead', 'tier3-worker', 'tier4-qa'], 
        required=True, 
        help="The tier role to execute"
    )
    parser.add_argument(
        "prompt", 
        type=str, 
        help="The prompt for the tier"
    )
    return parser

def main():
    parser = create_parser()
    args = parser.parse_args()
    
    docs = get_role_documents(args.role)
    print(f"Executing role: {args.role} with docs: {docs}")
    result = execute_agent(args.role, args.prompt, docs)
    print(result)

if __name__ == "__main__":
    main()