"""
SSH-based Firmware Flasher for OpenWrt Devices

Replicates the flash.ps1 logic using Python paramiko/scp:
  1. Auto-accept SSH host key
  2. Set device password via `passwd` command
  3. Upload firmware via SCP to /tmp/
  4. Verify, sync, and flash via sysupgrade
"""

import os
import time
import paramiko
from scp import SCPClient


def _create_ssh_client(ip, user, password, timeout=10):
    """Create an SSH client with auto-accept host key policy."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(ip, username=user, password=password, timeout=timeout,
                   look_for_keys=False, allow_agent=False)
    return client


import telnetlib

def set_device_password(ip, user="root", old_password="", new_password="admin123a",
                        status_cb=None):
    """
    Set device password via Telnet (if raw/reset) or SSH.
    """
    if status_cb:
        status_cb("Checking Telnet port for raw device...")

    # 1. Thử Telnet trước (OpenWrt mặc định mở Telnet 23 và cấm SSH Root khi chưa có Pass)
    try:
        tn = telnetlib.Telnet(ip, timeout=5)
        # Nếu vô được Telnet tức là thiết bị vừa Reset cứng chưa có pass
        if status_cb:
            status_cb("Telnet connected! Setting password...")
        
        # Đợi logo OpenWrt và prompt "root@OpenWrt:/# "
        time.sleep(1)
        tn.read_very_eager()
        
        # Gửi lệnh đổi pass
        tn.write(b"passwd\n")
        time.sleep(1)
        tn.write(new_password.encode('ascii') + b"\n")
        time.sleep(1)
        tn.write(new_password.encode('ascii') + b"\n")
        time.sleep(1)
        
        # Thoát telnet
        tn.write(b"exit\n")
        time.sleep(0.5)
        tn.close()
        
        # Chờ 3 giây để OpenWrt kịp đóng Telnet và nổ tiến trình Dropbear (SSH Server)
        if status_cb:
            status_cb("Password set via Telnet. Waiting for SSH to start...")
        time.sleep(3)
        return "DONE"
        
    except ConnectionRefusedError:
        # Port 23 đóng -> Tức là thiết bị đã có Pass và đã bật SSH, chuyển qua luồng mồi mật khẩu cũ
        pass
    except Exception as e:
        # Các lỗi timeout khác có thể do ping không tới
        return f"FAIL: Telnet check -> {e}"

    # 2. Rơi xuống luồng SSH nếu thiết bị cũ (cổng Telnet đóng)
    if status_cb:
        status_cb("Connecting SSH for password update...")

    try:
        client = _create_ssh_client(ip, user, old_password, timeout=5)
    except Exception as e:
        return f"FAIL: Cannot connect SSH (Old pass incorrect?) — {e}"

    try:
        if status_cb:
            status_cb("Setting password via SSH...")

        shell = client.invoke_shell()
        time.sleep(1)
        if shell.recv_ready():
            shell.recv(65535)

        shell.send("passwd\n")
        time.sleep(2)
        shell.send(f"{new_password}\n")
        time.sleep(1)
        shell.send(f"{new_password}\n")
        time.sleep(2)

        if shell.recv_ready():
            shell.recv(65535)

        shell.send("exit\n")
        time.sleep(0.5)
        shell.close()

        if status_cb:
            status_cb("Password set ✓")
        return "DONE"

    except Exception as e:
        return f"FAIL: {e}"
    finally:
        client.close()


def flash_device_ssh(ip, firmware_path, user="root", password="admin123a",
                     set_passwd=False, status_cb=None):
    """
    Flash firmware to an OpenWrt device via SSH/SCP.

    Steps (mirroring flash.ps1):
      1. (Optional) Set device password via passwd
      2. Upload firmware via SCP to /tmp/
      3. Verify uploaded file
      4. Sync filesystem
      5. Execute sysupgrade

    Returns:
        "DONE" on success, "FAIL: reason" on error
    """

    # ═══════════════════════════════════════════
    #  STEP 0: Set password (optional)
    # ═══════════════════════════════════════════
    if set_passwd:
        result = set_device_password(ip, user, "", password, status_cb)
        if result.startswith("FAIL"):
            # Try with current password as old password
            result = set_device_password(ip, user, password, password, status_cb)
            if result.startswith("FAIL"):
                return result

    # ═══════════════════════════════════════════
    #  STEP 1: Connect SSH
    # ═══════════════════════════════════════════
    if status_cb:
        status_cb("Connecting SSH...")

    try:
        client = _create_ssh_client(ip, user, password)
    except paramiko.AuthenticationException:
        return "FAIL: SSH authentication failed"
    except paramiko.SSHException as e:
        return f"FAIL: SSH error — {e}"
    except Exception as e:
        return f"FAIL: Cannot connect — {e}"

    try:
        # ═══════════════════════════════════════════
        #  STEP 2: Upload firmware via SCP
        # ═══════════════════════════════════════════
        if status_cb:
            status_cb("Uploading firmware via SCP...")

        filename = os.path.basename(firmware_path)
        remote_path = f"/tmp/{filename}"

        try:
            scp_client = SCPClient(client.get_transport(), socket_timeout=300)
            scp_client.put(firmware_path, remote_path)
            scp_client.close()
        except Exception as e:
            return f"FAIL: SCP upload failed — {e}"

        time.sleep(2)

        # ═══════════════════════════════════════════
        #  STEP 3: Verify firmware uploaded
        # ═══════════════════════════════════════════
        if status_cb:
            status_cb("Verifying firmware...")

        stdin, stdout, stderr = client.exec_command(
            f"test -f {remote_path} && ls -lh {remote_path}",
            timeout=10
        )
        verify_output = stdout.read().decode("utf-8", errors="ignore").strip()
        verify_err = stderr.read().decode("utf-8", errors="ignore").strip()

        if not verify_output:
            return f"FAIL: Firmware file not found on device after upload"

        # ═══════════════════════════════════════════
        #  STEP 4: Sync filesystem
        # ═══════════════════════════════════════════
        if status_cb:
            status_cb("Syncing filesystem...")

        client.exec_command("sync", timeout=10)
        time.sleep(2)

        # ═══════════════════════════════════════════
        #  STEP 5: Flash firmware (sysupgrade)
        # ═══════════════════════════════════════════
        if status_cb:
            status_cb("Flashing firmware (sysupgrade)...")

        try:
            # Capture output by redirecting to a file in /tmp first, or read from stdout
            # Use -F to force upgrade and bypass "Image metadata not present" error for uImage files
            stdin, stdout, stderr = client.exec_command(f"sysupgrade -F -v -n {remote_path} > /tmp/sysup.log 2>&1")
            
            # Wait up to 4 seconds to see if it immediately fails
            # Sysupgrade takes time and normally drops the connection, so if it finishes in < 4s, it's an error.
            time.sleep(4)
            
            if stdout.channel.exit_status_ready():
                exit_code = stdout.channel.recv_exit_status()
                
                # Fetch the error log
                _, log_out, _ = client.exec_command("cat /tmp/sysup.log")
                err_msg = log_out.read().decode("utf-8", errors="ignore").strip()
                
                return f"FAIL: sysupgrade terminated early (Code {exit_code}). Details:\n{err_msg}"

        except Exception:
            # Connection drop during sysupgrade is exactly what we expect if it succeeds
            pass

        if status_cb:
            status_cb("Rebooting...")
        time.sleep(3)

        return "DONE"

    except Exception as e:
        return f"FAIL: {e}"
    finally:
        try:
            client.close()
        except Exception:
            pass