"""Trigger and control OpsAgent workflows from local machine."""
from __future__ import annotations

import asyncio
import os
import sys
import argparse

from temporalio.client import Client

async def connect_client():
    address = os.environ.get("TEMPORAL_ADDRESS", "localhost:7233")
    namespace = os.environ.get("TEMPORAL_NAMESPACE", "default")
    print(f"Connecting to Temporal at {address}...")
    return await Client.connect(address, namespace=namespace)

async def start_workflow(args):
    client = await connect_client()
    wf_id = f"ops-{args.target_type}-{args.target_id.replace('/', '-')}"
    
    print(f"Starting OpsRemediateWorkflow for {args.target_type} target '{args.target_id}'...")
    try:
        handle = await client.start_workflow(
            "OpsRemediateWorkflow",
            args.target_type,
            args.target_id,
            id=wf_id,
            task_queue="ops",
        )
        print(f"Workflow started successfully! Workflow ID: {handle.id}")
        print("You can now monitor it, or run 'approve' / 'reject' command to signal it.")
    except Exception as e:
        print(f"Failed to start workflow: {e}")

async def send_signal(args, signal_type: str):
    client = await connect_client()
    wf_id = args.workflow_id
    
    print(f"Sending '{signal_type}' signal to workflow {wf_id}...")
    try:
        handle = client.get_workflow_handle(wf_id)
        if signal_type == "approve":
            await handle.signal("approve")
        else:
            await handle.signal("reject")
        print(f"Signal '{signal_type}' sent successfully!")
    except Exception as e:
        print(f"Failed to send signal: {e}")

async def get_result(args):
    client = await connect_client()
    wf_id = args.workflow_id
    
    print(f"Waiting for results of workflow {wf_id}...")
    try:
        handle = client.get_workflow_handle(wf_id)
        result = await handle.result()
        print("\n=== Workflow Final Result ===")
        import json
        print(json.dumps(result, indent=2))
    except Exception as e:
        print(f"Failed to retrieve workflow result: {e}")

def main():
    parser = argparse.ArgumentParser(description="Trigger and control OpsAgent workflows.")
    subparsers = parser.add_subparsers(dest="command", required=True)
    
    # Start command
    parser_start = subparsers.add_parser("start", help="Start a diagnostics & remediation workflow.")
    parser_start.add_argument("--target-type", required=True, choices=["kubernetes", "proxmox"])
    parser_start.add_argument("--target-id", required=True, help="e.g. default/web-pod or VMID like 101")
    
    # Approve command
    parser_app = subparsers.add_parser("approve", help="Send approve signal to running workflow.")
    parser_app.add_argument("--workflow-id", required=True)
    
    # Reject command
    parser_rej = subparsers.add_parser("reject", help="Send reject signal to running workflow.")
    parser_rej.add_argument("--workflow-id", required=True)
    
    # Result command
    parser_res = subparsers.add_parser("result", help="Wait for and display workflow final result.")
    parser_res.add_argument("--workflow-id", required=True)
    
    args = parser.parse_args()
    
    if args.command == "start":
        asyncio.run(start_workflow(args))
    elif args.command == "approve":
        asyncio.run(send_signal(args, "approve"))
    elif args.command == "reject":
        asyncio.run(send_signal(args, "reject"))
    elif args.command == "result":
        asyncio.run(get_result(args))

if __name__ == "__main__":
    main()