import os
import sys
import json
import platform
import subprocess
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"]

    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.binary_path = self._resolve_binary()
        
        # Cancellation support
        self._cancelled = threading.Event()
        self._current_process = None
        self._original_sigint_handler = None

    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}")

        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 not found_path:
            raise FileNotFoundError(
                f"Binary not found in '{self.binary_dir}'. \n"
                f"Expected one of: {candidates}. \n"
                f"Please compile the C++ tool and place it there."
            )

        # Ensure execution permissions on Linux/macOS
        if self.os_name in ["Linux", "Darwin"]:
            st = os.stat(found_path)
            os.chmod(found_path, st.st_mode | stat.S_IEXEC)

        return found_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
        
        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).
                            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()])
        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 self._current_process and self._current_process.poll() is None:
            self.cancel()
        self._restore_signal_handlers()
