import asyncio
import argparse
import os
import sys
import json

async def main():
    parser = argparse.ArgumentParser(description="Run OpsAgent diagnostics CLI tool.")
    parser.add_argument("--target-type", required=True, choices=["kubernetes", "proxmox"], help="Target platform to check.")
    parser.add_argument("--target-id", required=True, help="Specific target identifier (e.g., namespace/pod-name or Proxmox container VMID).")
    parser.add_argument("--approve", action="store_true", help="Approve and execute proposed remediation.")
    parser.add_argument("--real", action="store_true", help="Run real CLI commands instead of mocks.")
    
    args = parser.parse_args()
    
    # Configure logging
    import logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    
    # Set mock mode
    if args.real:
        os.environ["MOCK_OPS"] = "false"
        print("Executing in REAL mode...")
    else:
        os.environ["MOCK_OPS"] = "true"
        print("Executing in MOCK mode...")

    # We need to set mock key environment variables if they aren't present
    if "LITELLM_API_KEY" not in os.environ:
        os.environ["LITELLM_API_KEY"] = "mock-key"
    if "LITELLM_BASE_URL" not in os.environ:
        os.environ["LITELLM_BASE_URL"] = "http://localhost:4000" # fallback

    from ops_agent.graph import run_ops_graph
    
    print(f"Starting OpsAgent check for target type: {args.target_type}, ID: {args.target_id}...")
    
    try:
        result = await run_ops_graph(
            target_type=args.target_type,
            target_id=args.target_id,
            is_approved=args.approve
        )
        
        print("\n==================================================")
        print("                    RESULTS                       ")
        print("==================================================")
        print(f"Target Check: {result['target_type'].upper()} ({result['target_id']})")
        print("\n--- DIAGNOSIS ---")
        print(result["diagnosis"])
        print("\n--- PROPOSED ACTION ---")
        print(f"Action: {result['proposed_action']}")
        print(f"Details: {result['proposed_action_details']}")
        print("\n--- EXECUTION RESULT ---")
        print(result["action_result"])
        print("==================================================")
        
    except Exception as e:
        print(f"\nExecution failed: {str(e)}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())
