#!/usr/bin/python3

import os
import subprocess
import logging
import time
import argparse

skip_list = ["indi_eval", "indi_getprop", "indi_setprop", "indi_hid_test"]

# Driver source types
class DriverSource:
    """Driver source types."""
    SYSTEM = "system"
    FLATPAK_STABLE = "flatpak-stable"
    FLATPAK_MASTER = "flatpak-master"


# KStars flatpak app ID
KSTARS_APP_ID = "org.kde.kstars"

logging.basicConfig(format='%(asctime)s - %(levelname)s: %(message)s', level=logging.INFO)


def get_flatpak_driver_list(branch):
    """Get list of INDI drivers from flatpak installation."""
    try:
        # List drivers inside the flatpak sandbox at /app/bin
        cmd = [
            "flatpak", "run", "--branch=" + branch,
            "--command=sh", KSTARS_APP_ID,
            "-c", "ls /app/bin"
        ]
        output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode("utf-8")
        drivers = [f for f in output.splitlines() if f.startswith("indi_") and f not in skip_list]
        return drivers
    except subprocess.CalledProcessError as e:
        logging.error(f"Failed to get driver list from flatpak ({branch}): {e}")
        return []
    except FileNotFoundError:
        logging.error("Flatpak is not installed or KStars flatpak not found")
        return []


def get_system_driver_list():
    """Get list of INDI drivers from system installation."""
    drivers = []
    for file in os.listdir("/usr/bin"):
        if file.startswith("indi_") and file not in skip_list:
            drivers.append(file)
    return drivers


def build_indiserver_command(driver_source, driver_path, branch=None):
    """Build the indiserver command based on driver source."""
    if driver_source == DriverSource.SYSTEM:
        return f"indiserver -r 0 -v {driver_path}"
    elif driver_source in [DriverSource.FLATPAK_STABLE, DriverSource.FLATPAK_MASTER]:
        branch = branch or ("stable" if driver_source == DriverSource.FLATPAK_STABLE else "master")
        return (
            f"flatpak run --branch={branch} "
            f"--filesystem=/tmp --filesystem=~/.indi "
            f"--command=indiserver {KSTARS_APP_ID} "
            f"-r 0 -v {driver_path}"
        )
    else:
        raise ValueError(f"Invalid driver source: {driver_source}")


def test_driver(driver, command):
    """Test a single driver."""
    try:
        logging.info("------------------------------------------------------------------------------------------")
        os.system("pkill indiserver")
        time.sleep(2)
        output = subprocess.check_output(command, shell=True, timeout=2).decode("utf-8")
        if "Terminated" in output:
            logging.error("Driver failed: " + driver)
            return False
        return True
    except subprocess.CalledProcessError as error:
        logging.error("Driver failed: " + driver)
        return False
    except subprocess.TimeoutExpired:
        logging.info("%s Finished successfully" % driver)
        return True


def main():
    parser = argparse.ArgumentParser(
        description="Test INDI drivers from different sources"
    )
    parser.add_argument(
        "--driver-source",
        choices=[DriverSource.SYSTEM, DriverSource.FLATPAK_STABLE, DriverSource.FLATPAK_MASTER],
        default=DriverSource.SYSTEM,
        help="Source to run drivers from: 'system' (default), 'flatpak-stable', or 'flatpak-master'"
    )
    args = parser.parse_args()

    driver_source = args.driver_source
    logging.info(f"Testing drivers from source: {driver_source}")

    # Kill any existing indiserver processes
    os.system("pkill indiserver")

    # Get driver list based on source
    if driver_source == DriverSource.SYSTEM:
        drivers = get_system_driver_list()
        driver_path_template = "/usr/bin/{driver}"
        branch = None
    else:
        branch = "stable" if driver_source == DriverSource.FLATPAK_STABLE else "master"
        drivers = get_flatpak_driver_list(branch)
        driver_path_template = "/app/bin/{driver}"

    if not drivers:
        logging.error("No drivers found to test")
        exit(-1)

    logging.info(f"Found {len(drivers)} drivers to test")

    # Test each driver
    failed_drivers = []
    for driver in drivers:
        driver_path = driver_path_template.format(driver=driver)
        command = build_indiserver_command(driver_source, driver_path, branch)
        
        if not test_driver(driver, command):
            failed_drivers.append(driver)

    # Final cleanup
    os.system("pkill indiserver")

    # Report results
    logging.info("==========================================================================================")
    if failed_drivers:
        logging.error(f"Failed drivers ({len(failed_drivers)}): {', '.join(failed_drivers)}")
        exit(-1)
    else:
        logging.info(f"All {len(drivers)} drivers passed successfully")


if __name__ == "__main__":
    main()