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

# 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()
    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 create_tar(source_dir, tar_path):
    with tarfile.open(tar_path, "w:gz") as tar:
        tar.add(source_dir, arcname=os.path.basename(source_dir))

def main():
    parser = argparse.ArgumentParser(description="Upload modified Diagnostic to 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="SSH port")
    parser.add_argument("--local_dir", default="./Diagnostic", help="Local Diagnostic folder path")
    args = parser.parse_args()

    local_dir = args.local_dir
    if not os.path.exists(local_dir):
        cprint(f"[!] Local folder {local_dir} does not exist.", Color.RED)
        return

    local_tar = os.path.join(os.getcwd(), "Diagnostic_upload.tar.gz")
    cprint("[*] Creating tar of local Diagnostic folder...", Color.CYAN)
    create_tar(local_dir, local_tar)

    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:
        remote_extract_dir = "/home/pi/Extract"
        remote_tar_path = f"{remote_extract_dir}/Diagnostic_upload.tar.gz"
        remote_diag_dir = f"{remote_extract_dir}/Diagnostic"

        # Backup existing remote database folders
        cprint("[*] Creating backups of existing remote database folders...", Color.CYAN)
        out, err, status = run_command(ssh, "cat /home/pi/projects/imageSourceConfig.json")
        if status != 0 or err:
            cprint(f"[!] Failed to fetch imageSourceConfig.json: {err}", Color.RED)
            return
        try:
            config = json.loads(out)
        except json.JSONDecodeError:
            cprint("[!] Could not parse JSON, raw output:", Color.RED)
            cprint(out, Color.YELLOW)
            config = {}

        for key, target_path in config.items():
            backup_path = f"{target_path}_bak"
            # remove old backup
            run_command(ssh, f"sudo rm -rf {backup_path}", get_output=False)
            # move current database folder to backup
            run_command(ssh, f"if [ -d {target_path} ]; then sudo mv {target_path} {backup_path}; fi", get_output=False)
            cprint(f"[*] Backed up {target_path} -> {backup_path}", Color.GREEN)

        # Remove old Diagnostic folder if exists
        cprint(f"[*] Removing old {remote_diag_dir} if exists...", Color.YELLOW)
        run_command(ssh, f"sudo rm -rf {remote_diag_dir}", get_output=False)

        # Upload tar
        cprint("[*] Uploading Diagnostic tar to Pi...", Color.CYAN)
        try:
            with SCPClient(ssh.get_transport()) as scp:
                scp.put(local_tar, remote_tar_path)
        except Exception as e:
            cprint(f"[!] SCP upload failed: {e}", Color.RED)
        cprint(f"[*] Tar uploaded to {remote_tar_path}", Color.GREEN)

        # # Extract tar
        cprint(f"[*] Extracting tar {remote_tar_path} into {remote_extract_dir}", Color.CYAN)
        run_command(ssh, f"sudo tar -xzf {remote_tar_path} -C {remote_extract_dir}", get_output=False)
        run_command(ssh, f"sudo rm -f {remote_tar_path}", get_output=False)

        # Move each Diagnostic/<key> to target folder
        for key, target_path in config.items():
            source_path = f"{remote_diag_dir}/{key}"
            cprint(f"[*] Moving {source_path} -> {target_path}", Color.GREEN)
            run_command(ssh, f"sudo cp -r {source_path} {target_path}", get_output=False)

        cprint("[*] Replacement complete. Database backups done, Diagnostic copied.", Color.GREEN)

    finally:
        # Refresh raspicam
        cprint("[*] Refreshing raspicam...", Color.CYAN)
        run_command(ssh, "curl -s -o /dev/null http://localhost/pi-cgi/refresh.cgi", get_output=False)
        
        # Remove temporary Diagnostic folder on remote
        cprint(f"[*] Removing temporary {remote_diag_dir} folder on remote...", Color.YELLOW)
        run_command(ssh, f"sudo rm -rf {remote_diag_dir}", get_output=False)


        ssh.close()
        cprint("[*] SSH connection closed.", Color.GREEN)

        # Remove local tar
        if os.path.exists(local_tar):
            os.remove(local_tar)

if __name__ == "__main__":
    main()
