import argparse
import time
import paramiko
import json
import sys
import os
import tarfile
from scp import SCPClient

# ANSI color codes
class Color:
    RESET = "\033[0m"
    GREEN = "\033[92m"
    RED = "\033[91m"
    YELLOW = "\033[93m"
    CYAN = "\033[96m"

def cprint(text, color=Color.RESET, end="\n"):
    print(f"{color}{text}{Color.RESET}", end=end)

def run_command(ssh, command, get_output=True, decode=True):
    stdin, stdout, stderr = ssh.exec_command(command)
    exit_status = stdout.channel.recv_exit_status()  # wait for command to finish
    out = stdout.read()
    err = stderr.read()

    if get_output:
        if decode:
            return out.decode(errors="ignore").strip(), err.decode(errors="ignore").strip(), exit_status
        else:
            return out, err, exit_status
    return None, None, exit_status

def download_and_extract(local_tar, extract_to):
    with tarfile.open(local_tar, "r:gz") as tar_ref:
        tar_ref.extractall(extract_to)
    os.remove(local_tar)

def countdown(seconds):
    for i in range(seconds, 0, -1):
        print(f"\r[*] Waiting {i} seconds for raspicam to start up...", end="")
        time.sleep(1)
    print()  # newline after countdown

def main():
    parser = argparse.ArgumentParser(description="SSH control for raspicam")
    parser.add_argument("--ip", required=True, help="Target device IP")
    parser.add_argument("--user", required=True, help="SSH username")
    parser.add_argument("--password", required=True, help="SSH password")
    parser.add_argument("--port", type=int, default=22, help="use for frp" )
    parser.add_argument("--sleep", type=int, default=20, help="Wait time after raspicam start")
    parser.add_argument("--mode", type=int, default=0, help="0 is base raspicam, 1 is Extract raspicam")
    args = parser.parse_args()

    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(args.ip, port=args.port, username=args.user, password=args.password)

    try:
        if args.mode == 1:
            # Only stop/start/countdown in mode 1
            cprint("[*] Stopping raspicam...", Color.CYAN)
            _, _, _ = run_command(ssh, "sudo /etc/init.d/raspicam stop", get_output=False)

            cprint("[*] Starting raspicam...", Color.CYAN)
            _, _, _ = run_command(ssh, "nohup sudo /home/pi/Extract/raspicam > /dev/null 2>&1 &", get_output=False)

            countdown(args.sleep)
        else:
            cprint("[*] Mode 0 selected, skipping raspicam restart.", Color.YELLOW)

        cprint("[*] Fetching imageSourceConfig.json...", Color.CYAN)
        out, err, status = run_command(ssh, "cat /home/pi/projects/imageSourceConfig.json")
        if status != 0 or err:
            cprint(f"[!] Error reading config: {err}", Color.RED)
            config = {}
        else:
            try:
                config = json.loads(out)
                cprint(json.dumps(config, indent=4), Color.GREEN)
            except json.JSONDecodeError:
                cprint("[!] Could not parse JSON, raw output:", Color.RED)
                cprint(out, Color.YELLOW)
                config = {}

        # Prepare Diagnostic
        cprint("[*] Preparing Diagnostic folder...", Color.CYAN)
        _, _, _ = run_command(ssh, "sudo rm -rf /home/pi/Extract/Diagnostic && mkdir -p /home/pi/Extract/Diagnostic", get_output=False)

        for key, src_path in config.items():
            dest = f"/home/pi/Extract/Diagnostic/{key}"
            cmd = f"mkdir -p {dest} && cp -r {src_path}/* {dest}/"
            cprint(f"[*] Copying {src_path} -> {dest}", Color.YELLOW)
            _, _, _ = run_command(ssh, cmd, get_output=False)

        # Make Images folder
        _, _, _ = run_command(ssh, "sudo mkdir -p /home/pi/Extract/Diagnostic/Images", get_output=False)
        _, _, _ = run_command(ssh, "sudo mkdir -p /home/pi/Extract/Diagnostic/Images/Raw", get_output=False)
        _, _, _ = run_command(ssh, "sudo mkdir -p /home/pi/Extract/Diagnostic/Images/Rectified", get_output=False)
        cprint("[*] Diagnostic setup complete.", Color.GREEN)

        # Capture
        for key in config.keys():
            url = f"http://localhost/pi-cgi/captureRawLive.cgi?{key}=/home/pi/Extract/Diagnostic/Images/Raw"
            cprint(f"[*] Capturing for source {key} via {url}", Color.CYAN)
            _, _, status = run_command(ssh, f"curl -s -o /dev/null '{url}'", get_output=False)
            if status != 0:
                cprint(f"[!] Capture failed for source {key} (exit {status})", Color.RED)
            else:
                cprint(f"[*] Capture for source {key} completed successfully.", Color.GREEN)

        for key in config.keys():
            url = f"http://localhost/pi-cgi/captureRectified.cgi?{key}=/home/pi/Extract/Diagnostic/Images/Rectified"
            cprint(f"[*] Capturing for source {key} via {url}", Color.CYAN)
            _, _, status = run_command(ssh, f"curl -s -o /dev/null '{url}'", get_output=False)
            if status != 0:
                cprint(f"[!] Capture failed for source {key} (exit {status})", Color.RED)
            else:
                cprint(f"[*] Capture for source {key} completed successfully.", Color.GREEN)

        # Tar Diagnostic
        cprint("[*] Archiving Diagnostic folder with tar...", Color.CYAN)
        _, err, status = run_command(ssh, "cd /home/pi/Extract && sudo tar -czf Diagnostic.tar.gz Diagnostic")
        if status != 0:
            cprint(f"[!] Failed to create Diagnostic.tar.gz: {err}", Color.RED)
            return

        # Download tar
        local_tar = os.path.join(os.getcwd(), "Diagnostic.tar.gz")
        with SCPClient(ssh.get_transport()) as scp:
            scp.get("/home/pi/Extract/Diagnostic.tar.gz", local_tar)

        cprint("[*] Extracting Diagnostic.tar.gz locally...", Color.CYAN)
        download_and_extract(local_tar, os.getcwd())

        # Cleanup remote
        _, _, _ = run_command(ssh, "sudo rm -rf /home/pi/Extract/Diagnostic /home/pi/Extract/Diagnostic.tar.gz", get_output=False)
        cprint("[*] Cleanup done on remote.", Color.GREEN)

    finally:
        if args.mode == 1:
            cprint("[*] Restarting raspicam (cleanup)...", Color.CYAN)
            _, _, _ = run_command(ssh, "sudo /etc/init.d/raspicam restart", get_output=False)
        
        ssh.close()
        cprint("[*] Done.", Color.GREEN)

if __name__ == "__main__":
    main()
