# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "python-dotenv",
#     "questionary",
# ]
# ///

import os
import sys
import subprocess
import shutil
from pathlib import Path
from dotenv import load_dotenv

# Ensure we run from the script's directory
script_dir = Path(__file__).resolve().parent
os.chdir(script_dir)

try:
    import questionary
except ImportError:
    print("[Error] Failed to import questionary. Run this script using 'uv run flash.py'.")
    sys.exit(1)

def run_esphome(args, env):
    esphome_bin = shutil.which("esphome")
    if not esphome_bin:
        print("\n[Error] 'esphome' CLI not found on PATH.")
        print("Please install it (e.g. 'brew install esphome' or 'pip install esphome').\n")
        sys.exit(1)
        
    cmd = [esphome_bin] + args
    print(f"\nExecuting: {' '.join(cmd)}\n")
    try:
        process = subprocess.Popen(
            cmd,
            env=env,
            stdout=sys.stdout,
            stderr=sys.stderr,
            text=True
        )
        process.wait()
        return process.returncode
    except KeyboardInterrupt:
        print("\n\nExecution interrupted.")
        return 1

def get_device_name(yaml_path):
    try:
        with open(yaml_path, "r") as f:
            in_esphome = False
            for line in f:
                stripped = line.strip()
                if stripped.startswith("esphome:"):
                    in_esphome = True
                    continue
                if in_esphome:
                    if line.strip() and not line.startswith(" ") and not line.startswith("\t"):
                        break
                    if stripped.startswith("name:"):
                        name = stripped.split(":", 1)[1].split("#", 1)[0].strip()
                        return name.replace('"', '').replace("'", "")
    except Exception:
        pass
    return None

def main():
    env_file = script_dir / ".env"
    template_file = script_dir / ".env.template"

    # Auto-copy template if .env is missing
    if not env_file.exists():
        if template_file.exists():
            print(f"'.env' not found. Copying '.env.template' to '.env'...")
            shutil.copy(template_file, env_file)
            print("\n[Action Required] Please open '.env' and fill in your Wi-Fi credentials.")
            sys.exit(0)
        else:
            print("[Error] Neither '.env' nor '.env.template' could be found.")
            sys.exit(1)

    load_dotenv(env_file)
    
    # 1. Select YAML Configuration
    yaml_files = sorted([f.name for f in script_dir.glob("*.yaml")])
    if not yaml_files:
        print("[Error] No ESPHome configuration files (*.yaml) found in this directory.")
        sys.exit(1)
        
    selected_yaml = questionary.select(
        "Select the ESPHome configuration to compile/flash:",
        choices=yaml_files
    ).ask()
    
    if not selected_yaml:
        sys.exit(0)

    # 2. Select Command Action
    commands = [
        questionary.Choice("Run (Validate, compile, upload, and start logs)", "run"),
        questionary.Choice("Compile (Build only)", "compile"),
        questionary.Choice("Upload (Flash compiled binary only)", "upload"),
        questionary.Choice("Logs (Stream device logs)", "logs"),
        questionary.Choice("Validate (Check syntax only)", "config"),
        questionary.Choice("Clean (Clear build artifacts)", "clean"),
    ]
    
    selected_cmd = questionary.select(
        "Select ESPHome action to perform:",
        choices=commands
    ).ask()
    
    if not selected_cmd:
        sys.exit(0)

    # If upload is selected, verify build files exist
    if selected_cmd == "upload":
        device_name = get_device_name(script_dir / selected_yaml)
        if device_name:
            build_dir = script_dir / ".esphome" / "build" / device_name
            if not build_dir.exists():
                print(f"\n[Warning] Build directory for '{device_name}' does not exist.")
                print("You must compile the project before uploading it.")
                confirm = questionary.confirm(
                    "Would you like to compile + upload the configuration now (Run)?",
                    default=True
                ).ask()
                if confirm:
                    selected_cmd = "run"
                else:
                    sys.exit(0)

    # Prepare system environment
    env = os.environ.copy()
    
    # Execute ESPHome command
    exit_code = run_esphome([selected_cmd, selected_yaml], env)
    sys.exit(exit_code)

if __name__ == "__main__":
    main()