#!/usr/bin/env python3
"""
QEMU ISO Visual Screenshot Harvester

Features:
- Uses script directory as source
- Finds loose ISO files
- Finds ISO files inside ZIP archives
- Opens QEMU visibly (user watches boot)
- Sends keyboard test: print 2+2 ENTER
- Captures screenshot
- Saves all images to ./Screenshots
- Names screenshots after source ISO / ZIP
- Logs progress
- Standard library only
"""

import os
import sys
import time
import socket
import shutil
import zipfile
import tempfile
import subprocess
from pathlib import Path


SCRIPT_DIR = Path(__file__).resolve().parent
SCREENSHOT_DIR = SCRIPT_DIR / "Screenshots"
LOG_FILE = SCRIPT_DIR / "iso_screenshot_log.txt"

QEMU_CANDIDATES = [
    r"C:\Program Files\qemu\qemu-system-x86_64.exe",
    r"C:\Program Files\qemu\qemu-system-x86\_64.exe",
    "/usr/bin/qemu-system-x86_64",
    "/usr/local/bin/qemu-system-x86_64"
]

RAM_MB = 2048
BOOT_WAIT = 90
KEY_WAIT = 5
MONITOR_PORT = 4444


SCREENSHOT_DIR.mkdir(exist_ok=True)


def log(msg):
    print(msg, flush=True)
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(msg + "\n")


def find_qemu():

    for q in QEMU_CANDIDATES:
        if Path(q).exists():
            return q

    q = shutil.which("qemu-system-x86_64")

    return q



def clean_name(text):

    return "".join(
        c if c.isalnum() or c in "-_"
        else "_"
        for c in text
    )



def connect_monitor(port):

    for _ in range(60):

        try:
            return socket.create_connection(
                ("127.0.0.1", port),
                timeout=3
            )

        except OSError:
            time.sleep(1)

    return None



def monitor_commands(port, commands):

    sock = connect_monitor(port)

    if not sock:
        log("Monitor connection failed")
        return False

    try:

        stream = sock.makefile("w")

        for cmd in commands:
            stream.write(cmd + "\n")

        stream.flush()

        time.sleep(3)

        stream.close()
        sock.close()

        return True

    except Exception as e:

        log(f"Monitor error: {e}")

        return False



def keyboard_test_commands():

    return [

        "sendkey p",
        "sendkey r",
        "sendkey i",
        "sendkey n",
        "sendkey t",
        "sendkey spc",
        "sendkey 2",
        "sendkey kp_add",
        "sendkey 2",
        "sendkey ret"

    ]



def run_qemu(iso, output, port):

    log(f"Starting QEMU: {iso.name}")

    cmd = [

        QEMU,

        "-boot",
        "d",

        "-cdrom",
        str(iso),

        "-m",
        str(RAM_MB),

        "-vga",
        "std",

        "-monitor",
        f"tcp:127.0.0.1:{port},server,nowait"

    ]


    try:

        # Visible QEMU window
        proc = subprocess.Popen(cmd)

    except Exception as e:

        log(f"QEMU failed: {e}")

        return False


    log(
        f"Watching boot for {BOOT_WAIT} seconds..."
    )

    time.sleep(BOOT_WAIT)


    log(
        "Sending keyboard test: print 2+2"
    )


    commands = keyboard_test_commands()


    time.sleep(KEY_WAIT)


    commands.extend(
        [
            f"screendump {str(output).replace(chr(92), '/')}",
            "quit"
        ]
    )


    monitor_commands(
        port,
        commands
    )


    try:

        proc.wait(
            timeout=20
        )

    except subprocess.TimeoutExpired:

        log(
            "Force killing QEMU"
        )

        proc.kill()


    return output.exists()



def find_loose_isos():

    result = []

    for item in SCRIPT_DIR.iterdir():

        if item.is_file():

            if item.suffix.lower() == ".iso":

                result.append(
                    (
                        item,
                        clean_name(item.stem)
                    )
                )

    return result



def find_zip_isos():

    result = []

    for archive in SCRIPT_DIR.glob("*.zip"):

        log(
            f"Scanning ZIP: {archive.name}"
        )

        try:

            with zipfile.ZipFile(
                archive
            ) as z:

                for member in z.namelist():

                    if member.lower().endswith(".iso"):

                        result.append(
                            (
                                archive,
                                member
                            )
                        )

        except zipfile.BadZipFile:

            log(
                f"Bad ZIP: {archive}"
            )

    return result



def extract_iso(archive, member, temp):

    target = temp / Path(member).name

    with zipfile.ZipFile(
        archive
    ) as z:

        with z.open(member) as src:

            with open(target, "wb") as dst:

                shutil.copyfileobj(
                    src,
                    dst
                )

    return target



def main():

    global QEMU

    QEMU = find_qemu()

    if not QEMU:

        print(
            "QEMU not found"
        )

        sys.exit(1)


    open(
        LOG_FILE,
        "w",
        encoding="utf-8"
    ).close()


    jobs = []


    for iso, name in find_loose_isos():

        jobs.append(
            (
                "iso",
                iso,
                name
            )
        )


    for archive, member in find_zip_isos():

        jobs.append(
            (
                "zip",
                archive,
                member
            )
        )


    temp = Path(
        tempfile.mkdtemp(
            prefix="iso_temp_"
        )
    )


    total = len(jobs)


    log(
        f"Found {total} ISO images"
    )


    try:

        for index, job in enumerate(jobs):

            log("")
            log(
                f"[{index+1}/{total}]"
            )


            port = MONITOR_PORT + index


            if job[0] == "iso":

                iso = job[1]
                name = job[2]


            else:

                archive = job[1]
                member = job[2]

                log(
                    f"Extracting {member}"
                )

                iso = extract_iso(
                    archive,
                    member,
                    temp
                )

                name = clean_name(
                    archive.stem
                    + "_"
                    + Path(member).stem
                )


            output = (
                SCREENSHOT_DIR
                /
                f"{name}.png"
            )


            if output.exists():

                log(
                    f"Already exists: {output.name}"
                )

                continue


            if run_qemu(
                iso,
                output,
                port
            ):

                log(
                    f"Saved: {output}"
                )

            else:

                log(
                    f"Failed: {name}"
                )


    finally:

        shutil.rmtree(
            temp,
            ignore_errors=True
        )


    log(
        "COMPLETE"
    )



if __name__ == "__main__":
    main()