#!/usr/bin/env python3
"""
Raspberry Pi Setup Script for DReader Hardware.

This interactive script helps configure your DReader e-reader hardware by:
1. Detecting connected hardware (I2C devices, SPI, etc.)
2. Creating/editing hardware_config.json
3. Installing required system packages
4. Setting up permissions and services

Usage:
    sudo python3 setup_rpi.py
"""

import sys
import os
import json
import subprocess
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# Check if running on Raspberry Pi
try:
    with open('/proc/device-tree/model', 'r') as f:
        model = f.read()
        if 'Raspberry Pi' not in model:
            print("⚠️  Warning: This doesn't appear to be a Raspberry Pi")
            print(f"   Detected: {model.strip()}")
            response = input("Continue anyway? (y/N): ")
            if response.lower() != 'y':
                sys.exit(1)
except:
    print("⚠️  Warning: Could not detect Raspberry Pi")


class Colors:
    """ANSI color codes for terminal output."""
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    CYAN = '\033[96m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    END = '\033[0m'
    BOLD = '\033[1m'


def print_header(text: str):
    """Print a header."""
    print(f"\n{Colors.BOLD}{Colors.BLUE}{'='*70}{Colors.END}")
    print(f"{Colors.BOLD}{Colors.BLUE}{text:^70}{Colors.END}")
    print(f"{Colors.BOLD}{Colors.BLUE}{'='*70}{Colors.END}\n")


def print_step(number: int, text: str):
    """Print a step number."""
    print(f"\n{Colors.BOLD}{Colors.CYAN}Step {number}: {text}{Colors.END}")
    print(f"{Colors.CYAN}{'-'*70}{Colors.END}")


def print_success(text: str):
    """Print success message."""
    print(f"{Colors.GREEN}✓ {text}{Colors.END}")


def print_warning(text: str):
    """Print warning message."""
    print(f"{Colors.YELLOW}⚠ {text}{Colors.END}")


def print_error(text: str):
    """Print error message."""
    print(f"{Colors.RED}✗ {text}{Colors.END}")


def run_command(cmd: str, check: bool = True) -> Tuple[int, str, str]:
    """Run a shell command and return result."""
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            check=check
        )
        return result.returncode, result.stdout, result.stderr
    except subprocess.CalledProcessError as e:
        return e.returncode, e.stdout, e.stderr


def check_interfaces() -> Dict[str, bool]:
    """Check if I2C and SPI interfaces are enabled."""
    print("Checking system interfaces...")

    interfaces = {
        'i2c': False,
        'spi': False,
    }

    # Check I2C
    if os.path.exists('/dev/i2c-1'):
        interfaces['i2c'] = True
        print_success("I2C interface enabled")
    else:
        print_warning("I2C interface not enabled")

    # Check SPI
    if os.path.exists('/dev/spidev0.0'):
        interfaces['spi'] = True
        print_success("SPI interface enabled")
    else:
        print_warning("SPI interface not enabled")

    return interfaces


def detect_i2c_devices() -> List[str]:
    """Detect I2C devices."""
    print("\nScanning I2C bus...")

    returncode, stdout, stderr = run_command("i2cdetect -y 1", check=False)

    if returncode != 0:
        print_warning("Could not scan I2C bus (i2cdetect not found or no permission)")
        return []

    # Parse i2cdetect output
    devices = []
    for line in stdout.split('\n'):
        if ':' in line:
            # Extract hex addresses
            parts = line.split(':')[1].split()
            for part in parts:
                if part != '--' and len(part) == 2:
                    devices.append(f"0x{part}")

    if devices:
        print_success(f"Found {len(devices)} I2C device(s): {', '.join(devices)}")

        # Identify known devices
        device_map = {
            '0x38': 'FT5316 Touch Panel',
            '0x14': 'BMA400 Accelerometer',
            '0x15': 'BMA400 Accelerometer (alt)',
            '0x68': 'PCF8523 RTC',
            '0x40': 'INA219 Power Monitor',
        }

        print("\nDetected devices:")
        for addr in devices:
            device_name = device_map.get(addr, 'Unknown device')
            print(f"  {addr}: {device_name}")
    else:
        print_warning("No I2C devices detected")

    return devices


def enable_interfaces():
    """Enable I2C and SPI interfaces."""
    print("\nEnabling interfaces...")

    # Use raspi-config to enable I2C and SPI
    print("Enabling I2C...")
    run_command("raspi-config nonint do_i2c 0", check=False)

    print("Enabling SPI...")
    run_command("raspi-config nonint do_spi 0", check=False)

    print_success("Interfaces enabled (reboot required to take effect)")


def setup_permissions():
    """Set up user permissions for GPIO, I2C, and SPI."""
    print("\nSetting up user permissions...")

    user = os.environ.get('SUDO_USER', os.environ.get('USER'))

    groups = ['gpio', 'i2c', 'spi']
    for group in groups:
        print(f"Adding user '{user}' to group '{group}'...")
        returncode, _, _ = run_command(f"usermod -a -G {group} {user}", check=False)

        if returncode == 0:
            print_success(f"Added to {group} group")
        else:
            print_warning(f"Could not add to {group} group (may not exist)")

    print_warning("You must log out and back in for group changes to take effect")


def get_vcom_voltage() -> float:
    """Prompt user for VCOM voltage."""
    print("\n" + Colors.BOLD + "VCOM Voltage Configuration" + Colors.END)
    print("="*70)
    print("Your e-ink display has a VCOM voltage printed on a label.")
    print("This is usually on the back of the display.")
    print("")
    print("Example labels:")
    print("  • VCOM = -2.06V")
    print("  • VCOM: -1.98V")
    print("  • -2.14V")
    print("")
    print(Colors.RED + Colors.BOLD + "⚠️  IMPORTANT: Using incorrect VCOM can damage your display!" + Colors.END)
    print("")

    while True:
        vcom_str = input("Enter your display's VCOM voltage (e.g., -2.06): ").strip()

        try:
            vcom = float(vcom_str)

            if vcom > 0:
                print_warning("VCOM is usually negative. Did you forget the minus sign?")
                continue

            if vcom < -3.0 or vcom > -1.0:
                print_warning(f"VCOM {vcom}V is unusual. Most displays are between -1.5V and -2.5V")
                confirm = input("Are you sure this is correct? (y/N): ")
                if confirm.lower() != 'y':
                    continue

            return vcom

        except ValueError:
            print_error("Invalid voltage. Please enter a number (e.g., -2.06)")


def configure_gpio_buttons() -> dict:
    """Configure GPIO buttons interactively."""
    print("\n" + Colors.BOLD + "GPIO Button Configuration" + Colors.END)
    print("="*70)
    print("Configure physical buttons for navigation.")
    print("Buttons should be connected between GPIO pin and GND.")
    print("")

    enable = input("Enable GPIO buttons? (Y/n): ").strip().lower()
    if enable == 'n':
        return {
            "enabled": False,
            "pull_up": True,
            "bounce_time_ms": 200,
            "buttons": []
        }

    buttons = []

    # Common button configurations (based on actual hardware)
    button_presets = [
        ("prev_page", "Previous Page", "swipe_right", 22),
        ("next_page", "Next Page", "swipe_left", 27),
        ("power_off", "Power Off", "long_press", 21),
    ]

    print("\nAvailable GPIOs (BCM numbering): 2-27 (avoid 2, 3 if using I2C)")
    print("")

    for name, description, default_gesture, default_gpio in button_presets:
        print(f"\n{Colors.BOLD}{description} Button{Colors.END}")
        enable_btn = input(f"  Enable {description} button? (Y/n): ").strip().lower()

        if enable_btn == 'n':
            continue

        # Get GPIO pin
        while True:
            gpio_str = input(f"  GPIO pin (default {default_gpio}): ").strip()
            if not gpio_str:
                gpio = default_gpio
                break

            try:
                gpio = int(gpio_str)
                if gpio < 2 or gpio > 27:
                    print_error("    GPIO must be between 2 and 27")
                    continue
                if gpio in [2, 3]:
                    print_warning("    GPIO 2/3 are I2C pins (SDA/SCL)")
                    confirm = input("    Use anyway? (y/N): ")
                    if confirm.lower() != 'y':
                        continue
                break
            except ValueError:
                print_error("    Invalid GPIO number")

        # Add button
        buttons.append({
            "name": name,
            "gpio": gpio,
            "gesture": default_gesture,
            "description": description
        })

        print_success(f"  Configured: GPIO {gpio} -> {description}")

    return {
        "enabled": True,
        "pull_up": True,
        "bounce_time_ms": 200,
        "buttons": buttons
    }


def create_hardware_config(vcom: float, gpio_config: dict, i2c_devices: List[str]) -> dict:
    """Create hardware configuration dictionary."""
    # Auto-detect which optional components are available
    has_touch = '0x38' in i2c_devices
    has_accel = '0x14' in i2c_devices or '0x15' in i2c_devices
    has_rtc = '0x68' in i2c_devices
    has_power = '0x40' in i2c_devices

    config = {
        "_description": "Hardware configuration for DReader e-ink device",
        "_generated": "Generated by setup_rpi.py",

        "display": {
            "width": 1872,
            "height": 1404,
            "vcom": vcom,
            "spi_hz": 24000000,
            "auto_sleep": True
        },

        "gpio_buttons": gpio_config,

        "accelerometer": {
            "enabled": has_accel,
            "tilt_enabled": False,
            "orientation_enabled": has_accel,
            "calibration_file": "accelerometer_config.json"
        },

        "rtc": {
            "enabled": has_rtc
        },

        "power_monitor": {
            "enabled": has_power,
            "shunt_ohms": 0.1,
            "battery_capacity_mah": 3000,
            "low_battery_threshold": 20.0,
            "show_battery_interval": 100
        },

        "application": {
            "library_path": "/home/pi/Books",
            "auto_save_interval": 60,
            "force_library_mode": False,
            "log_level": "INFO"
        }
    }

    return config


def main():
    """Main setup function."""
    print_header("DReader Raspberry Pi Hardware Setup")

    # Check if running as root
    if os.geteuid() != 0:
        print_error("This script must be run with sudo")
        print("Usage: sudo python3 setup_rpi.py")
        sys.exit(1)

    # Step 1: Check interfaces
    print_step(1, "Checking System Interfaces")
    interfaces = check_interfaces()

    if not all(interfaces.values()):
        print("\nSome interfaces are not enabled.")
        enable = input("Enable I2C and SPI now? (Y/n): ").strip().lower()

        if enable != 'n':
            enable_interfaces()
            print_warning("Reboot required for interface changes to take effect")

    # Step 2: Detect hardware
    print_step(2, "Detecting I2C Devices")
    i2c_devices = detect_i2c_devices()

    if not i2c_devices:
        print_warning("No I2C devices detected. Check your wiring.")
        print("See HARDWARE_SETUP.md for wiring instructions.")

    # Step 3: Set up permissions
    print_step(3, "Setting Up User Permissions")
    setup_permissions()

    # Step 4: Configure VCOM
    print_step(4, "Display Configuration")
    vcom = get_vcom_voltage()
    print_success(f"VCOM voltage set to {vcom}V")

    # Step 5: Configure GPIO buttons
    print_step(5, "GPIO Button Configuration")
    gpio_config = configure_gpio_buttons()

    if gpio_config["enabled"]:
        print_success(f"Configured {len(gpio_config['buttons'])} button(s)")
    else:
        print("GPIO buttons disabled")

    # Step 6: Generate configuration
    print_step(6, "Generating Configuration File")
    config = create_hardware_config(vcom, gpio_config, i2c_devices)

    config_file = Path("hardware_config.json")
    with open(config_file, 'w') as f:
        json.dump(config, f, indent=2)

    print_success(f"Configuration saved to {config_file}")

    # Step 7: Summary
    print_header("Setup Complete!")

    print("Configuration summary:")
    print(f"  • Display: {config['display']['width']}x{config['display']['height']}, VCOM={config['display']['vcom']}V")
    print(f"  • GPIO Buttons: {'Enabled' if gpio_config['enabled'] else 'Disabled'}")
    if gpio_config['enabled']:
        for btn in gpio_config['buttons']:
            print(f"      - {btn['description']}: GPIO {btn['gpio']}")
    print(f"  • Accelerometer: {'Enabled' if config['accelerometer']['enabled'] else 'Disabled'}")
    print(f"  • RTC: {'Enabled' if config['rtc']['enabled'] else 'Disabled'}")
    print(f"  • Power Monitor: {'Enabled' if config['power_monitor']['enabled'] else 'Disabled'}")

    print("\n" + Colors.BOLD + "Next Steps:" + Colors.END)
    print("1. Review and edit hardware_config.json if needed")
    print("2. Reboot if you enabled I2C/SPI: sudo reboot")
    print("3. Log out and back in for permission changes")
    print("4. Run: python examples/run_on_hardware_config.py")
    print("")


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nSetup cancelled by user")
        sys.exit(1)
    except Exception as e:
        print_error(f"Setup failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)