#!/usr/bin/env python3
"""
Universal SCPI Interface - Supports both Serial (USB) and WiFi/TCP connections

Provides a unified interface that automatically detects and connects via:
1. WiFi/TCP (if available) - preferred for logging use cases
2. USB Serial - fallback when WiFi not available

Features:
- Automatic WiFi IP discovery via USB connection
- Seamless fallback between connection types
- Same interface as SCPISerial for drop-in replacement
- Connection status monitoring
- OPC synchronization support

Usage:
    from scpi_universal import SCPIUniversal

    # Auto-detect (tries WiFi, then USB)
    with SCPIUniversal() as instr:
        print(instr.query("*IDN?"))
        print(f"Connected via: {instr.connection_type}")

    # Force WiFi with specific IP
    with SCPIUniversal(prefer_wifi=True, wifi_host="192.168.1.100") as instr:
        temp = instr.query(":SENS:TEMP?")

    # Force USB only
    with SCPIUniversal(prefer_wifi=False) as instr:
        status = instr.query(":PID:STAT?")

(c) 2026 Prof. Flavio ABREU ARAUJO. All rights reserved.
"""

import contextlib
import socket
import time
from typing import Optional, Union

from .scpi_serial import SCPISerial

# Default SCPI TCP port (standard)
DEFAULT_SCPI_PORT = 5025


class TCPSocket:
    """
    TCP socket wrapper that mimics serial.Serial interface for SCPI over WiFi.

    This allows using the same command/query interface for both serial and TCP.
    """

    def __init__(self, host: str, port: int = DEFAULT_SCPI_PORT, timeout: float = 2.0):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.sock: Optional[socket.socket] = None
        self._buffer = b""
        self.port_name = f"TCP:{host}:{port}"  # For compatibility with serial interface

    def connect(self) -> None:
        """Connect to the TCP server."""
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.settimeout(self.timeout)
        try:
            self.sock.connect((self.host, self.port))
        except (OSError, socket.timeout) as e:
            self.sock = None
            raise ConnectionError(f"Failed to connect to {self.host}:{self.port}: {e}") from e

    def close(self) -> None:
        """Close the connection."""
        if self.sock:
            with contextlib.suppress(Exception):
                self.sock.close()
            self.sock = None

    def write(self, data: Union[bytes, str]) -> None:
        """Write data to socket (accepts bytes or string)."""
        if isinstance(data, str):
            data = data.encode("utf-8")
        if self.sock:
            self.sock.sendall(data)

    def readline(self) -> bytes:
        """Read a line from socket (returns bytes)."""
        if not self.sock:
            return b""

        while b"\n" not in self._buffer:
            try:
                chunk = self.sock.recv(4096)
                if not chunk:
                    break
                self._buffer += chunk
            except socket.timeout:
                break

        if b"\n" in self._buffer:
            line, self._buffer = self._buffer.split(b"\n", 1)
            return line + b"\n"
        else:
            # Return whatever we have if no newline
            line = self._buffer
            self._buffer = b""
            return line

    def reset_input_buffer(self) -> None:
        """Clear input buffer."""
        self._buffer = b""
        if self.sock:
            self.sock.setblocking(False)
            try:
                while True:
                    chunk = self.sock.recv(4096)
                    if not chunk:
                        break
            except (OSError, BlockingIOError):
                pass
            finally:
                self.sock.setblocking(True)
                self.sock.settimeout(self.timeout)

    def reset_output_buffer(self) -> None:
        """No-op for TCP (no output buffer to clear)."""
        pass

    def flush(self) -> None:
        """No-op for TCP (sends are immediate)."""
        pass

    @property
    def in_waiting(self) -> int:
        """Check if data is available."""
        if self._buffer:
            return len(self._buffer)
        if not self.sock:
            return 0

        # Non-blocking check for available data
        self.sock.setblocking(False)
        try:
            chunk = self.sock.recv(4096, socket.MSG_PEEK)
            return len(chunk)
        except (OSError, BlockingIOError):
            return 0
        finally:
            self.sock.setblocking(True)
            self.sock.settimeout(self.timeout)


class SCPIUniversal:
    """
    Universal SCPI interface supporting both Serial (USB) and WiFi/TCP connections.

    Automatically detects and connects via the best available method:
    1. WiFi/TCP (if enabled and available) - preferred for logging scenarios
    2. USB Serial - fallback when WiFi not available

    Provides same interface as SCPISerial for seamless integration.
    """

    def __init__(
        self,
        port: Optional[str] = None,
        wifi_host: Optional[str] = None,
        wifi_port: int = DEFAULT_SCPI_PORT,
        prefer_wifi: bool = True,
        baudrate: int = 115200,
        timeout: float = 1.0,
        auto_connect: bool = True,
        sync_mode: bool = True,
    ):
        """
        Initialize Universal SCPI interface.

        Args:
            port: Serial port path (e.g., "/dev/cu.usbmodem2101", "COM3")
                  If None, will attempt auto-detection
            wifi_host: WiFi IP address (e.g., "192.168.1.100")
                       If None, will attempt auto-detection via USB
            wifi_port: WiFi TCP port (default: 5025)
            prefer_wifi: Try WiFi connection first (default: True)
            baudrate: Serial baudrate (default: 115200)
            timeout: Connection timeout in seconds (default: 1.0)
            auto_connect: Automatically connect on initialization (default: True)
            sync_mode: Use OPC synchronization for write commands (default: True)
        """
        self.port = port
        self.wifi_host = wifi_host
        self.wifi_port = wifi_port
        self.prefer_wifi = prefer_wifi
        self.baudrate = baudrate
        self.timeout = timeout
        self.sync_mode = sync_mode

        # Internal connection objects
        self._ser: Optional[SCPISerial] = None
        self._tcp: Optional[TCPSocket] = None
        self._connected = False
        self.connection_type: str = "disconnected"  # "serial", "tcp", or "disconnected"

        if auto_connect:
            self.connect()

    @property
    def port_name(self) -> str:
        """Get the current connection identifier."""
        if self.connection_type == "tcp" and self._tcp:
            return self._tcp.port_name
        elif self.connection_type == "serial" and self._ser:
            return self._ser.port_name
        return "Not connected"

    @property
    def serial(self):
        """Get the underlying serial/TCP object for direct access."""
        if self.connection_type == "tcp":
            return self._tcp
        elif self.connection_type == "serial" and self._ser:
            return self._ser.serial
        return None

    def __enter__(self):
        """Context manager entry - returns connected instance."""
        if not self._connected:
            self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - closes connection."""
        self.close()
        return False

    def _discover_wifi_via_usb(self) -> Optional[str]:
        """
        Discover WiFi IP address by temporarily connecting via USB.

        Returns:
            IP address string or None if not available

        Note: Opens the port without exclusive locking. If another process
        (e.g., temp-logger-gui) already owns the port, both may read the
        same OS serial buffer simultaneously, causing minor cross-talk.
        This is accepted here because the discovery window is very short
        (~200 ms) and the result drives a switch to WiFi/TCP which avoids
        any ongoing USB contention.
        """
        temp_serial = None
        try:
            # Temporarily connect via USB to query WiFi status.
            # NOTE: Do NOT use exclusive=True here. On macOS, TIOCEXCL set during
            # this temporary open is NOT cleared when this fd is closed while another
            # process (e.g., temp-logger-gui) still holds the port open. That leaves
            # TIOCEXCL set on the device, blocking all subsequent port opens (including
            # the USB fallback in connect()) with EBUSY.
            temp_serial = SCPISerial(
                port=self.port,
                baudrate=self.baudrate,
                timeout=self.timeout,
                auto_connect=True,
                exclusive=False,  # must NOT hold exclusive lock during discovery
            )

            # Check if WiFi is connected
            time.sleep(0.1)  # Small delay for stability
            temp_serial.serial.reset_input_buffer()
            connected = temp_serial.query(":WIFI:CONNECTED?").strip()

            if connected not in ("1", "TRUE", "ON"):
                return None

            # Get IP address
            time.sleep(0.1)
            temp_serial.serial.reset_input_buffer()
            ip_addr = temp_serial.query(":WIFI:IP?").strip()

            # Validate IP address format
            if ip_addr and ip_addr != "0.0.0.0" and ip_addr != "None":
                parts = ip_addr.split(".")
                if len(parts) == 4 and all(p.isdigit() and 0 <= int(p) <= 255 for p in parts):
                    return ip_addr

            return None

        except Exception:
            return None
        finally:
            if temp_serial:
                temp_serial.close()

    def _try_tcp_connection(self, host: str, port: int) -> bool:
        """
        Attempt to establish TCP connection.

        Returns:
            True if successful, False otherwise
        """
        try:
            self._tcp = TCPSocket(host, port, timeout=self.timeout)
            self._tcp.connect()

            # Verify connection with simple query
            self._tcp.reset_input_buffer()
            self._tcp.write(b"*IDN?\n")
            response = self._tcp.readline()

            if response:
                self.connection_type = "tcp"
                self._connected = True
                return True
            else:
                self._tcp.close()
                self._tcp = None
                return False

        except Exception:
            if self._tcp:
                self._tcp.close()
                self._tcp = None
            return False

    def _try_serial_connection(self) -> bool:
        """
        Attempt to establish serial connection.

        Returns:
            True if successful, False otherwise
        """
        try:
            self._ser = SCPISerial(
                port=self.port,
                baudrate=self.baudrate,
                timeout=self.timeout,
                auto_connect=True,
            )
            self.connection_type = "serial"
            self._connected = True
            return True

        except Exception:
            if self._ser:
                self._ser.close()
                self._ser = None
            return False

    def connect(self) -> None:
        """
        Connect to instrument via best available method.

        Connection priority (if prefer_wifi=True):
        1. WiFi/TCP with provided host
        2. WiFi/TCP with auto-discovered host (via USB query)
        3. USB Serial  ← WARNING: can cause cross-talk if temp-logger-gui is running

        Connection priority (if prefer_wifi=False):
        1. USB Serial
        2. WiFi/TCP with provided host
        3. WiFi/TCP with auto-discovered host

        Raises:
            ConnectionError: If all connection methods fail
        """
        if self._connected:
            return

        if self.prefer_wifi:
            # Try WiFi first

            # 1. Explicit wifi_host provided — try TCP only, no USB fallback
            if self.wifi_host:
                if self._try_tcp_connection(self.wifi_host, self.wifi_port):
                    return
                raise ConnectionError(
                    f"Failed to connect to WiFi at {self.wifi_host}:{self.wifi_port}. "
                    "Ensure the device is powered on and connected to WiFi."
                )

            # 2. No explicit host — try auto-discovering WiFi IP via USB
            # NOTE: _discover_wifi_via_usb() uses exclusive port access so it fails
            # gracefully when another tool (e.g., temp-logger-gui) already owns the port.
            discovered_ip = self._discover_wifi_via_usb()
            if discovered_ip and self._try_tcp_connection(discovered_ip, self.wifi_port):
                self.wifi_host = discovered_ip  # Store for future reference
                return

            # 3. Fall back to USB serial
            # WARNING: If another tool (e.g., temp-logger-gui) is already connected
            # via USB serial, this will cause cross-talk: both tools read from and
            # write to the same serial port, corrupting each other's SCPI traffic.
            # To avoid this, always specify wifi_host explicitly when WiFi is available.
            if self._try_serial_connection():
                return

        else:
            # Try USB first

            # 1. Try USB serial
            if self._try_serial_connection():
                return

            # 2. Try provided WiFi host
            if self.wifi_host and self._try_tcp_connection(self.wifi_host, self.wifi_port):
                return

            # 3. Try auto-discovering WiFi IP (requires temporary USB connection)
            discovered_ip = self._discover_wifi_via_usb()
            if discovered_ip and self._try_tcp_connection(discovered_ip, self.wifi_port):
                self.wifi_host = discovered_ip
                return

        # All connection attempts failed
        raise ConnectionError(
            "Failed to connect via any method. "
            "Ensure device is powered on and either USB is connected or WiFi is enabled."
        )

    def close(self) -> None:
        """Close the active connection."""
        if self._tcp:
            self._tcp.close()
            self._tcp = None
        if self._ser:
            self._ser.close()
            self._ser = None
        self._connected = False
        self.connection_type = "disconnected"

    def write(self, command: str) -> None:
        """
        Write SCPI command to instrument.

        Args:
            command: SCPI command string (newline added automatically)
        """
        if not self._connected:
            raise ConnectionError("Not connected")

        if not command.endswith("\n"):
            command += "\n"

        if self.connection_type == "tcp" and self._tcp:
            self._tcp.write(command.encode("utf-8"))
        elif self.connection_type == "serial" and self._ser:
            self._ser.write(command)

        # OPC synchronization if enabled
        if self.sync_mode and not command.strip().endswith("?"):
            time.sleep(0.02)  # Small delay before OPC
            self.query("*OPC?")

    def query(self, command: str, timeout: Optional[float] = None) -> str:
        """
        Send SCPI query and read response.

        Args:
            command: SCPI query string (newline added automatically)
            timeout: Optional timeout override (only works for TCP connections)

        Returns:
            Response string (stripped of whitespace)
        """
        if not self._connected:
            raise ConnectionError("Not connected")

        if not command.endswith("\n"):
            command += "\n"

        if self.connection_type == "tcp" and self._tcp:
            # Save original timeout if override requested
            original_timeout = self._tcp.timeout
            if timeout is not None:
                self._tcp.sock.settimeout(timeout)

            try:
                self._tcp.reset_input_buffer()
                self._tcp.write(command.encode("utf-8"))
                response = self._tcp.readline().decode("utf-8", errors="ignore").strip()
                return response
            finally:
                if timeout is not None:
                    self._tcp.sock.settimeout(original_timeout)

        elif self.connection_type == "serial" and self._ser:
            # SCPISerial doesn't support timeout parameter in query()
            # It uses the timeout set during initialization
            return self._ser.query(command)

        return ""

    def query_binary(self, command: str, timeout: Optional[float] = None) -> bytes:
        """
        Send SCPI query and read binary response.

        Args:
            command: SCPI query string
            timeout: Optional timeout override (only works for TCP connections)

        Returns:
            Raw bytes response
        """
        if not self._connected:
            raise ConnectionError("Not connected")

        if not command.endswith("\n"):
            command += "\n"

        if self.connection_type == "tcp" and self._tcp:
            original_timeout = self._tcp.timeout
            if timeout is not None:
                self._tcp.sock.settimeout(timeout)

            try:
                self._tcp.reset_input_buffer()
                self._tcp.write(command.encode("utf-8"))
                return self._tcp.readline()
            finally:
                if timeout is not None:
                    self._tcp.sock.settimeout(original_timeout)

        elif self.connection_type == "serial" and self._ser:
            # SCPISerial doesn't have query_binary, use query and encode
            response = self._ser.query(command)
            return response.encode("utf-8")

        return b""


def main():
    """Test/demo for SCPIUniversal."""
    import sys

    print("SCPIUniversal Test Program")
    print("=" * 50)

    # Parse command line
    wifi_host = None
    prefer_wifi = True

    if len(sys.argv) > 1:
        if sys.argv[1] in ("-h", "--help"):
            print("Usage: python scpi_universal.py [wifi_ip]")
            print("  wifi_ip: Optional WiFi IP address (e.g., 192.168.1.100)")
            print("           If omitted, will auto-detect via USB")
            return 0
        wifi_host = sys.argv[1]

    try:
        print("\nConnecting to instrument...")
        with SCPIUniversal(wifi_host=wifi_host, prefer_wifi=prefer_wifi) as instr:
            print(f"✓ Connected via: {instr.connection_type.upper()}")
            print(f"  Port/Address: {instr.port_name}\n")

            # Test basic queries
            print("Testing basic queries:")
            print("-" * 50)

            idn = instr.query("*IDN?")
            print(f"*IDN?: {idn}")

            temp = instr.query(":SENS:TEMP?")
            print(f":SENS:TEMP?: {temp}")

            if instr.connection_type == "tcp":
                print("\n✓ WiFi/TCP connection successful!")
            else:
                print("\n✓ USB Serial connection successful!")

            print("\nConnection test passed!")

    except ConnectionError as e:
        print(f"\n❌ Connection failed: {e}")
        return 1
    except Exception as e:
        print(f"\n❌ Error: {e}")
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
