#! /usr/bin/env python3
import argparse
import sys
import logging
import requests
import pygame
import time
import subprocess

logger = logging.getLogger(__name__)

def main() -> int:
    # Handle program arguments
    ap = argparse.ArgumentParser(prog='pwnagotchi-watch', description='Watch the pwnagotchi screen remotely')
    ap.add_argument("--host", help="Pwnagotchi hostname or IP address", default="10.0.0.2")
    ap.add_argument("--port", help="Pwnagotchi port", type=int, default=8080)
    ap.add_argument("-u", "--username", help="Pwnagotchi username", default="changeme")
    ap.add_argument("-p", "--password", help="Pwnagotchi password", default="changeme")
    ap.add_argument("--refresh-rate", help="Refresh rate in seconds", type=int, default=2)
    ap.add_argument("--pin", help="Pin the window on top", action="store_true")
    ap.add_argument('-v', '--verbose', help='Enable verbose logging', action='store_true')
    args = ap.parse_args()

    # Configure logging
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format='%(levelname)s:	%(message)s',
    )
    
    # Create a window to draw the image to
    pygame.init()
    screen = pygame.display.set_mode((250,122))
    pygame.display.set_caption(f"pwnagotchi ({args.host})")
    
    # Pin the window on top
    if args.pin:
        logger.info("Pinning the window on top")
        window_id = pygame.display.get_wm_info()["window"]
        subprocess.run(["wmctrl", "-i", "-r", str(window_id), "-b", "add,above"], check=True)
    
    # Continuously fetch the pwnagotchi screen
    while True:
        
        # If the window is closed, exit
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return 0
        
        # Attempt to fetch the pwnagotchi screen
        try:
            # Fetch the pwnagotchi screen
            logger.debug(f"Fetching pwnagotchi screen from {args.host}:{args.port}")
            try:
                response = requests.get(f"http://{args.host}:{args.port}/ui", auth=(args.username, args.password), timeout=1)
                response.raise_for_status()
            except Exception as e: 
                logger.error(f"Failed to connect to pwnagotchi at {args.host}:{args.port}")
                logger.error(f"Error: {e}")
                time.sleep(args.refresh_rate)
                continue
            screen_data = response.content
            
            # Draw the screen to the window
            logger.debug("Drawing pwnagotchi screen")
            with open("/tmp/pwnscreen.png", "wb") as f:
                f.write(screen_data)
            screen.fill((0,0,0))
            try:
                img = pygame.image.load("/tmp/pwnscreen.png")
                screen.blit(img, (0,0))
                pygame.display.flip()
            except pygame.error as e:
                logger.error(f"Failed to load pwnagotchi screen: {e}")
            
        except requests.RequestException as e:
            logger.error(f"Failed to fetch pwnagotchi screen: {e}")
            
        time.sleep(args.refresh_rate)
        

    return 0

if __name__ == "__main__":
    sys.exit(main())