import os
import sys
import json
import platform
import subprocess
import stat
import shutil
import signal
import threading
import time
from pathlib import Path
from typing import Optional, List, Dict, Union

class K230FlashTool:
    """
    A Python wrapper for the Kendryte K230 C++ Flash CLI.
    Wraps the binary execution to handle OS detection, device listing, and flashing.
    """

    # Map OS name to binary filename
    BIN_MAP = {
        "Windows": "k230_flash_cli.exe",
        "Linux": "k230_flash_cli", 
        "Darwin": "k230_flash_cli"
    }

    # Valid storage media types defined in k230_flash.cpp enum
    VALID_MEDIA = ["EMMC", "SDCARD", "SPI_NAND", "SPI_NOR", "OTP"]
    DEFAULT_LINUX_FLASH_TOOL = Path("k230-flash")

    def __init__(self):
        """
        Initialize the wrapper.
        :param binary_dir: Path to the folder containing the compiled k230_flash binaries.
        """
        self.os_name = platform.system()
        self.binary_dir = Path(__file__).resolve().parent / "bin"
        self.cli_style = "k230_flash_cli"
        self.device_index = os.getenv("K230_FLASH_DEVICE", "0")
        self.use_sudo = False

        # Cancellation support
        self._cancelled = threading.Event()
        self._current_process = None
        self._original_sigint_handler = None

        self.binary_path = self._resolve_binary()

    def _style_for_binary(self, path: Path) -> str:
        if path.name == "k230-flash":
            return "k230-flash"
        return "k230_flash_cli"

    def _should_use_sudo(self, style: str) -> bool:
        env_value = os.getenv("K230_FLASH_USE_SUDO")
        if env_value is not None:
            return env_value.lower() not in ("0", "false", "no", "off")

        return False

    def _prepare_binary(self, path: Path) -> Path:
        if self.os_name in ["Linux", "Darwin"] and not os.access(path, os.X_OK):
            try:
                st = os.stat(path)
                os.chmod(path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
            except OSError:
                pass

        if self.os_name in ["Linux", "Darwin"] and not os.access(path, os.X_OK):
            raise PermissionError(f"Flash tool is not executable: {path}")

        self.cli_style = self._style_for_binary(path)
        self.use_sudo = self._should_use_sudo(self.cli_style)
        return path

    def _find_linux_k230_flash(self) -> Optional[Path]:
        candidates = []

        if self.DEFAULT_LINUX_FLASH_TOOL.is_absolute():
            candidates.append(self.DEFAULT_LINUX_FLASH_TOOL)

        path_binary = shutil.which("k230-flash")
        if path_binary:
            candidates.append(Path(path_binary))

        candidates.append(Path(sys.executable).resolve().parent / "k230-flash")

        try:
            import site
            candidates.append(Path(site.getuserbase()) / "bin" / "k230-flash")
        except (AttributeError, ImportError):
            pass

        for candidate in candidates:
            if candidate.is_file():
                return candidate

        return None

    def _install_linux_k230_flash(self) -> None:
        if os.getenv("K230_FLASH_AUTO_INSTALL", "1").lower() in ("0", "false", "no", "off"):
            return

        cmd = [sys.executable, "-m", "pip", "install", "k230-flash"]
        print(f"Linux k230-flash tool not found. Installing with: {' '.join(cmd)}")
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace"
        )

        if result.stdout:
            print(result.stdout, end="" if result.stdout.endswith("\n") else "\n")

        if result.returncode != 0:
            raise RuntimeError(f"Failed to install k230-flash with pip (exit status {result.returncode})")

    def _resolve_binary(self) -> Path:
        """Finds the correct binary for the current OS and ensures it is executable."""
        if self.os_name not in self.BIN_MAP:
            raise OSError(f"Unsupported OS: {self.os_name}")

        env_binary = os.getenv("K230_FLASH_TOOL")
        if env_binary:
            env_path = Path(env_binary).expanduser()
            if env_path.is_file():
                return self._prepare_binary(env_path)
            raise FileNotFoundError(f"K230_FLASH_TOOL does not point to a file: {env_path}")

        bin_name = self.BIN_MAP[self.os_name]
        
        # Search for exact match or generic 'k230_flash'
        candidates = [bin_name, "k230_flash", "k230_flash_cli"]

        found_path = None
        for c in candidates:
            p = self.binary_dir / c
            if p.exists():
                found_path = p
                break

        if found_path:
            return self._prepare_binary(found_path)

        if self.os_name == "Linux":
            external_binary = self._find_linux_k230_flash()
            if external_binary:
                return self._prepare_binary(external_binary)

            self._install_linux_k230_flash()
            external_binary = self._find_linux_k230_flash()
            if external_binary:
                return self._prepare_binary(external_binary)

        raise FileNotFoundError(
            f"Binary not found in '{self.binary_dir}'. \n"
            f"Expected one of: {candidates}. \n"
            f"Also checked K230_FLASH_TOOL and k230-flash on PATH. \n"
            f"Please install a Linux flash tool or set K230_FLASH_TOOL to its full path."
        )

    def cancel(self):
        """Cancel the current running process"""
        self._cancelled.set()
        if self._current_process:
            try:
                # Try graceful termination first
                self._current_process.terminate()
                
                # Wait a bit for graceful shutdown
                try:
                    self._current_process.wait(timeout=2)
                except subprocess.TimeoutExpired:
                    # Force kill if not responding
                    self._current_process.kill()
                    self._current_process.wait()
            except (ProcessLookupError, OSError):
                pass  # Process already terminated

    def _setup_signal_handlers(self):
        """Set up signal handlers for Ctrl+C"""
        def signal_handler(sig, frame):
            print("\nReceived interrupt signal, cancelling operation...")
            self.cancel()
            
        self._original_sigint_handler = signal.signal(signal.SIGINT, signal_handler)

    def _restore_signal_handlers(self):
        """Restore original signal handlers"""
        if self._original_sigint_handler:
            signal.signal(signal.SIGINT, self._original_sigint_handler)

    def _run(self, args: List[str], stream: bool = True, parse_json: bool = False, 
              timeout: Optional[int] = 10) -> Union[str, List, Dict]:
        """
        Executes the binary with arguments with proper cancellation and timeout support.

        Usage for Terminal Output + Result:
            output_str = self._run(["-some", "args"], stream=True, parse_json=False)
        
        :param args: List of arguments.
        :param stream: If True, print stdout in real-time (default behavior).
        :param parse_json: If True, capture output and attempt to parse as JSON.
        :param timeout: Optional timeout in seconds.
        :return: The output string (if parse_json=False) or Dict/List (if parse_json=True).
        """
        self._cancelled.clear()
        cmd = [str(self.binary_path)] + args
        if self.use_sudo:
            cmd.insert(0, "sudo")
        
        try:
            # Set up signal handlers for Ctrl+C
            self._setup_signal_handlers()

            # Use Popen to stream output
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                universal_newlines=True,
                encoding='utf-8',
                errors='replace'
            )
            self._current_process = process

            captured_lines = []
            start_time = time.time()

            # Read output with cancellation and timeout checks
            if process.stdout:
                while True:
                    # Check for cancellation
                    if self._cancelled.is_set():
                        process.terminate()
                        raise KeyboardInterrupt("Operation cancelled by user")
                    
                    # Check for timeout
                    if timeout and (time.time() - start_time) > timeout:
                        process.terminate()
                        raise TimeoutError(f"Command timed out after {timeout} seconds: {' '.join(cmd)}")
                    
                    # Try to read line with short timeout
                    try:
                        line = process.stdout.readline()
                        if not line:  # EOF reached
                            break

                        captured_lines.append(line)
                        if stream and not parse_json:
                            print(line, end='', flush=True)
                    except (IOError, ValueError):
                        break  # Pipe closed or other error

            # Wait for process completion with timeout
            try:
                remaining_time = None
                if timeout:
                    elapsed = time.time() - start_time
                    remaining_time = max(0, timeout - elapsed)
                
                process.wait(timeout=remaining_time)
            except subprocess.TimeoutExpired:
                process.terminate()
                try:
                    process.wait(timeout=2)  # Give it 2 seconds to terminate
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait()
                raise TimeoutError(f"Command timed out after {timeout} seconds: {' '.join(cmd)}")

            # Check if process was terminated due to cancellation
            if self._cancelled.is_set():
                raise KeyboardInterrupt("Operation cancelled by user")

            full_output = "".join(captured_lines)

            # Handle process return codes
            if process.returncode != 0:
                # Allow non-zero exit only if we are asking for help or listing devices
                # Also allow if process was terminated by signal
                if (process.returncode in [-2, -15] or  # SIGINT, SIGTERM
                    any(arg in args for arg in ["--help", "-h", "--list-device", "-l"])):
                    return full_output
                raise subprocess.CalledProcessError(process.returncode, cmd, output=full_output)

            if parse_json:
                try:
                    # Try to parse the entire output as JSON first
                    return json.loads(full_output.strip())
                except json.JSONDecodeError:
                    # If that fails, try to find JSON structure in the output
                    start = full_output.find('[')
                    end = full_output.rfind(']') + 1
                    if start != -1 and end != -1:
                        json_str = full_output[start:end]
                        try:
                            return json.loads(json_str)
                        except json.JSONDecodeError:
                            pass

                    # If JSON parsing fails, return raw text
                    if stream:
                        print("Warning: Expected JSON but got raw text.")
                    return full_output

            return full_output
            
        except FileNotFoundError:
            raise FileNotFoundError(f"Flash tool binary not found: {self.binary_path}")
        except PermissionError:
            raise PermissionError(f"Permission denied accessing flash tool binary: {self.binary_path}")
        finally:
            self._current_process = None
            self._restore_signal_handlers()

    def flash(self, 
              image_path: str, 
              medium_type: str = "SDCARD",
              auto_reboot: bool = True,
              timeout: Optional[int] = 10) -> None:
        """
        Flashes a firmware image to the K230.

        :param image_path: Path to the firmware file (e.g., sysimage-sdcard.img or sysimage_sdcard.kdimg).
                            Get this from list_devices().
        :param medium_type: Target medium: EMMC, SDCARD, SPI_NAND, SPI_NOR, or OTP.
        :param auto_reboot: If True, adds --auto-reboot flag.
        :param timeout: Timeout in seconds for flash operation (default: 300 seconds = 5 minutes)
        """
        img = Path(image_path)
        if not img.exists():
            raise FileNotFoundError(f"Image file not found: {image_path}")

        if medium_type.upper() not in self.VALID_MEDIA:
            raise ValueError(f"Invalid medium '{medium_type}'. Valid: {self.VALID_MEDIA}")

        # Construct arguments
        args = []
        args.extend(["-m", medium_type.upper()])
        if self.cli_style == "k230-flash":
            if auto_reboot:
                args.append("--auto-reboot")
            if img.suffix.lower() == ".kdimg":
                args.append(str(img.resolve()))
            else:
                args.extend([self.device_index, str(img.resolve())])
        else:
            args.extend(["-f", str(img.resolve())])
            if auto_reboot:
                args.append("--auto-reboot")

        print(f"--- Flashing {medium_type} with Command {args} ---")
        print("Press Ctrl+C to cancel the operation")
        
        try:
            self._run(args, stream=True, timeout=timeout)
            print("\n--- Flash Complete ---")
        except KeyboardInterrupt:
            print("\n--- Flash Operation Cancelled by User ---")
            raise
        except TimeoutError:
            print(f"\n--- Flash Operation Timed Out after {timeout} seconds ---")
            raise

    def __del__(self):
        """Ensure cleanup on object destruction"""
        if getattr(self, "_current_process", None) and self._current_process.poll() is None:
            self.cancel()
        if hasattr(self, "_original_sigint_handler"):
            self._restore_signal_handlers()
