import argparse
import os

from smolagents.monitoring import LogLevel

from screensuite.agents.models import QwenVLModel
from screensuite.agents.vision_agents.desktop_agent import DesktopAgent
from screensuite.agents.vision_agents.e2b_agent import E2BVisionAgent


def parse_args():
    parser = argparse.ArgumentParser(description="Run the E2B Qwen Vision Agent")
    parser.add_argument(
        "--task", type=str, default=None, help="Task to perform (if not provided, will use default example task)"
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=os.environ.get("E2B_API_KEY"),
        help="E2B API key (default: E2B_API_KEY environment variable)",
    )
    parser.add_argument(
        "--resolution", type=str, default="1024,768", help="Screen resolution as width,height (default: 1024,768)"
    )
    parser.add_argument(
        "--model-path",
        type=str,
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Path to Qwen2.5VL model (default: Qwen/Qwen2.5-VL-3B-Instruct)",
    )
    parser.add_argument(
        "--device", type=str, default="auto", help="Device to run model on: 'cuda', 'cpu', or 'auto' (default: auto)"
    )
    parser.add_argument("--max-steps", type=int, default=20, help="Maximum number of steps (default: 20)")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
    return parser.parse_args()


def run_agent_example(e2b_provider: bool = False):
    args = parse_args()

    # Validate E2B API key
    if not args.api_key:
        raise ValueError("E2B API key not provided. Set E2B_API_KEY environment variable or use --api-key")

    # Parse resolution
    try:
        width, height = map(int, args.resolution.split(","))
    except ValueError:
        raise ValueError("Invalid resolution format. Use 'width,height' (e.g., '1024,768')")

    print(f"Initializing with resolution: {width}x{height}")
    print(f"Using model: {args.model_path} on device: {args.device}")

    # Initialize the Qwen2.5VL model
    try:
        model = QwenVLModel(model_path=args.model_path, device=args.device)
        print("Model initialized successfully")
    except Exception as e:
        raise RuntimeError(f"Failed to initialize model: {str(e)}")

    # Initialize the E2B Vision Agent
    agent = None
    # try:
    if e2b_provider:
        agent = E2BVisionAgent(
            model=model,
            data_dir="./output",
            e2b_api_key=args.api_key,
            resolution=(width, height),
            max_steps=args.max_steps,
            verbosity_level=LogLevel.INFO,
        )
    else:
        agent = DesktopAgent(
            model=model,
            data_dir="./output",
            max_steps=args.max_steps,
            verbosity_level=LogLevel.INFO,
        )
    print("Agent initialized successfully")

    # Use provided task or default example task
    if args.task:
        task = args.task
    else:
        task = """Please follow these steps:
        1. Open Google (google.com) in the browser
        2. Search for "Python programming tutorial"
        3. Find and click on a Python tutorial link
        4. Scroll down to see the content
        5. Take a screenshot and tell me what topics are covered in the tutorial
        """
        # task = "show me a photo of a dog"

    print("\n=== Starting task execution ===")
    print(f"Task: {task}")
    print("================================\n")

    # Run the agent with the task
    result = agent.run(task)

    print("\n=== Task execution completed ===")
    print(f"Result: {result}")
    print("=================================")

    # except Exception as e:
    #     print(f"\n⚠️ Error during execution: {str(e)}")
    #     raise
    # finally:
    #     # Always close the agent to clean up resources
    #     if agent:
    #         print("\nClosing agent and cleaning up resources...")
    #         agent.close()
    #         print("Cleanup complete")


if __name__ == "__main__":
    try:
        run_agent_example()
    except Exception as e:
        print(f"Fatal error: {str(e)}")
        exit(1)
