import argparse
import csv
import json
import sys
import time
from datetime import datetime
from pathlib import Path

try:
    import pythoncom
    import win32com.client
    from win32com.client import VARIANT
except ImportError as exc:
    raise SystemExit(
        "Missing dependency: pywin32. Install with: python -m pip install --user pywin32"
    ) from exc


WMI_SCOPE = r"root\WMI"
WMI_CLASS = "MICommonInterface"
WMI_INSTANCE_NAME = r"ACPI\PNP0C14\MIFS_0"
WMI_METHOD = "MiInterface"

METHOD_GET = 0xFA
METHOD_SET = 0xFB
CMD_SYSTEM_PER_MODE = 8
CMD_GPU_MODE = 9
CMD_FN_LOCK = 11
CMD_TP_LOCK = 12
CMD_CPU_GPU_FAN_SPEED = 13
CMD_AMBIENT_LIGHT = 15
CMD_RGB_KEYBOARD_MODE = 16
CMD_RGB_KEYBOARD_BRIGHTNESS = 18
CMD_SYSTEM_AC_TYPE = 19
CMD_MAX_FAN_SPEED_SWITCH = 20
CMD_MAX_FAN_SPEED = 21
CMD_CPU_THERMOMETER = 22

SYSTEM_MODE_VALUES = {
    "balance": 0,
    "performance": 1,
    "quiet": 2,
}

SYSTEM_MODE_NAMES = {0: "Balance", 1: "Performance", 2: "Quiet"}
GPU_MODE_NAMES = {0: "Hybrid", 1: "Discrete"}
AC_TYPE_NAMES = {1: "TypeC", 2: "BarrelJack"}
ON_OFF_NAMES = {0: "off", 1: "on"}

RGB_MODE_NAMES = {
    0: "Off",
    1: "RGBAutoCyclic",
    2: "RGBFixedMode",
    3: "CustomColors",
}


class FanReader:
    def __init__(self, instance_name=None):
        locator = win32com.client.Dispatch("WbemScripting.SWbemLocator")
        self.service = locator.ConnectServer(".", WMI_SCOPE)
        self.device = self._find_device(instance_name or WMI_INSTANCE_NAME)
        self.method = self.device.Methods_(WMI_METHOD)

    def _find_device(self, instance_name):
        wanted = instance_name.lower()
        devices = list(self.service.ExecQuery(f"SELECT * FROM {WMI_CLASS}"))
        if not devices:
            raise RuntimeError(f"No {WMI_CLASS} instances found in {WMI_SCOPE}")

        for device in devices:
            current = str(device.Properties_("InstanceName").Value)
            if current.lower() == wanted:
                return device

        for device in devices:
            current = str(device.Properties_("InstanceName").Value)
            if wanted in current.lower() or current.lower() in wanted:
                return device

        names = ", ".join(str(d.Properties_("InstanceName").Value) for d in devices)
        raise RuntimeError(f"{WMI_CLASS} instance not found: {instance_name!r}; found: {names}")

    def _exec(self, method_type, command, values=None):
        payload = [0] * 32
        payload[1] = method_type
        payload[3] = command
        if values:
            for index, value in enumerate(values):
                payload[4 + index] = value

        in_params = self.method.InParameters.SpawnInstance_()
        in_params.Properties_("InData").Value = VARIANT(
            pythoncom.VT_ARRAY | pythoncom.VT_UI1,
            payload,
        )

        out_params = self.device.ExecMethod_(WMI_METHOD, in_params)
        out_data = out_params.Properties_("OutData").Value
        if out_data is None:
            raise RuntimeError("WMI returned empty OutData")
        return list(out_data)

    def read_cpu_gpu_rpm(self):
        data = self._exec(METHOD_GET, CMD_CPU_GPU_FAN_SPEED)
        if len(data) < 8:
            raise RuntimeError(f"WMI OutData too short: {data!r}")

        cpu_rpm = (data[5] << 8) + data[4]
        gpu_rpm = (data[7] << 8) + data[6]
        return cpu_rpm, gpu_rpm, data

    def read_byte(self, command):
        data = self._exec(METHOD_GET, command)
        if len(data) < 5:
            raise RuntimeError(f"WMI OutData too short: {data!r}")
        return data[4], data

    def read_cpu_temp(self):
        value, data = self.read_byte(CMD_CPU_THERMOMETER)
        return value, data

    def read_max_fan_switch(self):
        value, data = self.read_byte(CMD_MAX_FAN_SPEED_SWITCH)
        return value != 0, data

    def read_max_fan_speed(self):
        value, data = self.read_byte(CMD_MAX_FAN_SPEED)
        return value, data

    def set_value(self, command, value):
        data = self._exec(METHOD_SET, command, [value])
        return data

    def set_max_fan_switch(self, enabled):
        return self.set_value(CMD_MAX_FAN_SPEED_SWITCH, 1 if enabled else 0)

    def set_max_fan_speed(self, speed):
        return self.set_value(CMD_MAX_FAN_SPEED, speed)

    def set_system_mode(self, mode):
        return self.set_value(CMD_SYSTEM_PER_MODE, mode)

    def probe_get(self, command):
        return self._exec(METHOD_GET, command)

    def read_full_snapshot(self):
        cpu_rpm, gpu_rpm, _ = self.read_cpu_gpu_rpm()
        cpu_temp, _ = self.read_cpu_temp()
        fields = {}
        for label, command in [
            ("system_mode", CMD_SYSTEM_PER_MODE),
            ("gpu_mode", CMD_GPU_MODE),
            ("fn_lock", CMD_FN_LOCK),
            ("tp_lock", CMD_TP_LOCK),
            ("ambient", CMD_AMBIENT_LIGHT),
            ("rgb_brightness", CMD_RGB_KEYBOARD_BRIGHTNESS),
            ("rgb_mode", CMD_RGB_KEYBOARD_MODE),
            ("ac_type", CMD_SYSTEM_AC_TYPE),
            ("max_fan_switch", CMD_MAX_FAN_SPEED_SWITCH),
            ("max_fan_speed", CMD_MAX_FAN_SPEED),
        ]:
            try:
                value, _ = self.read_byte(command)
                fields[label] = value
            except Exception:
                fields[label] = None
        fields["cpu_rpm"] = cpu_rpm
        fields["gpu_rpm"] = gpu_rpm
        fields["cpu_temp"] = cpu_temp
        return fields


def enrich_snapshot(fields):
    def n(table, key):
        v = fields.get(key)
        return table.get(v, "") if v is not None else ""

    return {
        **fields,
        "system_mode_name": n(SYSTEM_MODE_NAMES, "system_mode"),
        "gpu_mode_name": n(GPU_MODE_NAMES, "gpu_mode"),
        "ac_type_name": n(AC_TYPE_NAMES, "ac_type"),
        "rgb_mode_name": n(RGB_MODE_NAMES, "rgb_mode"),
        "max_fan_switch_name": n(ON_OFF_NAMES, "max_fan_switch"),
        "fn_lock_name": n(ON_OFF_NAMES, "fn_lock"),
        "tp_lock_name": n(ON_OFF_NAMES, "tp_lock"),
        "ambient_name": n(ON_OFF_NAMES, "ambient"),
    }


FULL_CSV_FIELDNAMES = [
    "ts_iso",
    "cpu_rpm",
    "gpu_rpm",
    "cpu_temp",
    "system_mode",
    "system_mode_name",
    "gpu_mode",
    "gpu_mode_name",
    "max_fan_switch",
    "max_fan_switch_name",
    "max_fan_speed",
    "ac_type",
    "ac_type_name",
    "fn_lock",
    "fn_lock_name",
    "tp_lock",
    "tp_lock_name",
    "ambient",
    "ambient_name",
    "rgb_mode",
    "rgb_mode_name",
    "rgb_brightness",
    "guard",
]

MINIMAL_CSV_FIELDNAMES = ["ts_iso", "cpu_rpm", "gpu_rpm", "cpu_temp", "guard"]


def append_csv_row(path, fieldnames, row):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    new_file = not path.exists() or path.stat().st_size == 0
    safe = {k: ("" if v is None else v) for k, v in row.items()}
    with path.open("a", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
        if new_file:
            writer.writeheader()
        writer.writerow(safe)


def format_full_snapshot(fields):
    def name(table, value):
        return table.get(value, f"?{value}") if value is not None else "?"

    parts = [
        f"CPU Fan: {fields.get('cpu_rpm', '?'):>5} RPM",
        f"GPU Fan: {fields.get('gpu_rpm', '?'):>5} RPM",
        f"CPU Temp: {fields.get('cpu_temp', '?'):>3} C",
        f"Mode: {name(SYSTEM_MODE_NAMES, fields.get('system_mode'))}",
        f"GPU: {name(GPU_MODE_NAMES, fields.get('gpu_mode'))}",
        f"MaxFan: {name(ON_OFF_NAMES, fields.get('max_fan_switch'))}({fields.get('max_fan_speed', '?')})",
        f"AC: {name(AC_TYPE_NAMES, fields.get('ac_type'))}",
        f"Fn: {name(ON_OFF_NAMES, fields.get('fn_lock'))}",
        f"TP: {name(ON_OFF_NAMES, fields.get('tp_lock'))}",
        f"Ambient: {name(ON_OFF_NAMES, fields.get('ambient'))}",
        f"RGB: {name(RGB_MODE_NAMES, fields.get('rgb_mode'))} L{fields.get('rgb_brightness', '?')}",
    ]
    now = datetime.now().strftime("%H:%M:%S")
    return now + "  " + "  ".join(parts)


def format_line(cpu_rpm, gpu_rpm, cpu_temp=None, guard_state=None, raw=None):
    now = datetime.now().strftime("%H:%M:%S")
    line = f"{now}  CPU Fan: {cpu_rpm:5d} RPM  GPU Fan: {gpu_rpm:5d} RPM"
    if cpu_temp is not None:
        line += f"  CPU Temp: {cpu_temp:3d} C"
    if guard_state:
        line += f"  guard={guard_state}"
    if raw is not None:
        line += f"  raw={raw[:8]}"
    return line


def format_control_state(reader):
    switch, _ = reader.read_max_fan_switch()
    speed, _ = reader.read_max_fan_speed()
    return f"readback: max_fan={'on' if switch else 'off'}, speed={speed}"


def format_probe_result(command, data):
    raw = data[:16]
    byte4 = data[4] if len(data) > 4 else None
    word45 = ((data[5] << 8) + data[4]) if len(data) > 5 else None
    word67 = ((data[7] << 8) + data[6]) if len(data) > 7 else None
    return (
        f"cmd={command:3d} raw={raw} "
        f"byte4={byte4} word45={word45} word67={word67}"
    )


def main():
    parser = argparse.ArgumentParser(
        description="Realtime CPU/GPU fan RPM reader for MICommonInterface.",
    )
    parser.add_argument(
        "-i",
        "--interval",
        type=float,
        default=1.0,
        help="poll interval in seconds, default: 1.0",
    )
    parser.add_argument("--once", action="store_true", help="read once and exit")
    parser.add_argument("--raw", action="store_true", help="show first 8 raw OutData bytes")
    parser.add_argument(
        "--line",
        action="store_true",
        help="print one line per sample instead of updating the same line",
    )
    parser.add_argument(
        "--instance",
        default=WMI_INSTANCE_NAME,
        help=f"{WMI_CLASS} InstanceName, default: {WMI_INSTANCE_NAME}",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help=f"list {WMI_CLASS} instances and exit",
    )
    parser.add_argument(
        "--max-fan",
        choices=("on", "off"),
        help="turn max/manual fan mode on or off before reading",
    )
    parser.add_argument(
        "--set-speed",
        type=int,
        help="set manual fan speed byte value. Original UI range is usually 22..58.",
    )
    parser.add_argument(
        "--no-auto-enable",
        action="store_true",
        help="do not automatically enable max/manual fan mode before --set-speed",
    )
    parser.add_argument(
        "--show-temp",
        action="store_true",
        help="also read and display CPU temperature from WMI command 22",
    )
    parser.add_argument(
        "--guard-temp",
        type=int,
        help="enable guard mode: set max/manual fan when CPU temp is >= this value",
    )
    parser.add_argument(
        "--guard-release-temp",
        type=int,
        help="in guard mode, restore previous fan settings when CPU temp is <= this value",
    )
    parser.add_argument(
        "--guard-speed",
        type=int,
        default=58,
        help="fan speed byte used by guard mode, default: 58",
    )
    parser.add_argument(
        "--restore-speed",
        type=int,
        help="restore this speed when guard releases instead of the startup value",
    )
    parser.add_argument(
        "--system-mode",
        choices=tuple(SYSTEM_MODE_VALUES),
        help="set system performance mode: balance, performance, or quiet",
    )
    parser.add_argument(
        "--turbo",
        action="store_true",
        help="turbo preset: force performance mode, max/manual fan on, speed from "
        "--turbo-speed unless --set-speed is already set",
    )
    parser.add_argument(
        "--turbo-speed",
        type=int,
        default=58,
        help="with --turbo: fan level byte if --set-speed omitted (default 58; UI often ~22..58)",
    )
    parser.add_argument(
        "--full",
        action="store_true",
        help="display all available WMI fields each tick",
    )
    parser.add_argument(
        "--probe-get",
        nargs=2,
        type=int,
        metavar=("START", "END"),
        help="read-only probe: execute GET commands in inclusive range and exit",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="each tick print one UTF-8 JSON line; with --full includes all WMI snapshot fields",
    )
    parser.add_argument(
        "--log-csv",
        type=Path,
        metavar="PATH",
        help="append one CSV row per tick (with --full: all WMI columns; else rpm/temp/guard)",
    )
    args = parser.parse_args()

    if args.interval <= 0:
        raise SystemExit("interval must be > 0")
    if args.guard_temp is not None:
        if not 0 <= args.guard_temp <= 255:
            raise SystemExit("--guard-temp must be between 0 and 255")
        if not 0 <= args.guard_speed <= 255:
            raise SystemExit("--guard-speed must be between 0 and 255")
        if args.guard_release_temp is None:
            args.guard_release_temp = max(0, args.guard_temp - 5)
        if args.guard_release_temp > args.guard_temp:
            raise SystemExit("--guard-release-temp must be <= --guard-temp")
    if args.restore_speed is not None and not 0 <= args.restore_speed <= 255:
        raise SystemExit("--restore-speed must be between 0 and 255")
    if args.probe_get is not None:
        start, end = args.probe_get
        if not 0 <= start <= end <= 255:
            raise SystemExit("--probe-get requires 0 <= START <= END <= 255")
    if args.turbo:
        if args.system_mode is not None and args.system_mode != "performance":
            print("Note: --turbo forces system mode to performance.", file=sys.stderr)
        args.system_mode = "performance"
        args.max_fan = "on"
        if args.set_speed is None:
            if not 0 <= args.turbo_speed <= 255:
                raise SystemExit("--turbo-speed must be between 0 and 255")
            args.set_speed = args.turbo_speed
    if args.json and not args.full and not (args.show_temp or args.guard_temp is not None):
        print(
            "Note: --json without --full only includes rpm; add --show-temp for cpu_temp.",
            file=sys.stderr,
        )

    reader = None
    pythoncom.CoInitialize()
    try:
        if args.list:
            locator = win32com.client.Dispatch("WbemScripting.SWbemLocator")
            service = locator.ConnectServer(".", WMI_SCOPE)
            for device in service.ExecQuery(f"SELECT * FROM {WMI_CLASS}"):
                print(device.Properties_("InstanceName").Value)
            return

        reader = FanReader(args.instance)
        if args.probe_get is not None:
            start, end = args.probe_get
            for command in range(start, end + 1):
                try:
                    data = reader.probe_get(command)
                    print(format_probe_result(command, data), flush=True)
                except Exception as exc:
                    print(f"cmd={command:3d} error={exc}", flush=True)
            return

        previous_switch = None
        previous_speed = None
        guard_active = False

        if args.system_mode:
            mode_value = SYSTEM_MODE_VALUES[args.system_mode]
            reader.set_system_mode(mode_value)
            print(f"System mode: {args.system_mode} ({mode_value})", flush=True)

        if args.guard_temp is not None:
            try:
                previous_switch, _ = reader.read_max_fan_switch()
                previous_speed, _ = reader.read_max_fan_speed()
            except Exception:
                previous_switch = False
                previous_speed = None
            print(
                "Guard mode: "
                f"CPU >= {args.guard_temp} C -> speed {args.guard_speed}; "
                f"CPU <= {args.guard_release_temp} C -> restore",
                flush=True,
            )

        if args.max_fan:
            reader.set_max_fan_switch(args.max_fan == "on")
            print(
                f"Max/manual fan mode: {args.max_fan}; {format_control_state(reader)}",
                flush=True,
            )

        if args.set_speed is not None:
            if not 0 <= args.set_speed <= 255:
                raise SystemExit("--set-speed must be between 0 and 255")
            if not args.no_auto_enable:
                reader.set_max_fan_switch(True)
                print("Max/manual fan mode: on", flush=True)
            reader.set_max_fan_speed(args.set_speed)
            print(
                f"Manual fan speed value: {args.set_speed}; {format_control_state(reader)}",
                flush=True,
            )

        while True:
            if args.full:
                fields = reader.read_full_snapshot()
                cpu_rpm = fields["cpu_rpm"]
                gpu_rpm = fields["gpu_rpm"]
                cpu_temp = fields["cpu_temp"]
                raw = None
            else:
                cpu_rpm, gpu_rpm, raw = reader.read_cpu_gpu_rpm()
                cpu_temp = None
                fields = None

            guard_state = None

            if (args.show_temp or args.guard_temp is not None) and cpu_temp is None:
                cpu_temp, _ = reader.read_cpu_temp()

            if args.guard_temp is not None:
                if cpu_temp >= args.guard_temp and not guard_active:
                    reader.set_max_fan_switch(True)
                    reader.set_max_fan_speed(args.guard_speed)
                    guard_active = True
                    print(
                        f"\nGuard trigger: CPU {cpu_temp} C >= {args.guard_temp} C, "
                        f"set speed {args.guard_speed}; {format_control_state(reader)}",
                        flush=True,
                    )
                elif cpu_temp <= args.guard_release_temp and guard_active:
                    restore_speed = (
                        args.restore_speed
                        if args.restore_speed is not None
                        else previous_speed
                    )
                    if restore_speed is not None:
                        reader.set_max_fan_speed(restore_speed)
                    reader.set_max_fan_switch(bool(previous_switch))
                    guard_active = False
                    print(
                        f"\nGuard release: CPU {cpu_temp} C <= {args.guard_release_temp} C, "
                        f"restore fan settings; {format_control_state(reader)}",
                        flush=True,
                    )
                guard_state = "max" if guard_active else "normal"

            if args.full and fields is not None:
                line = format_full_snapshot(fields)
                if guard_state:
                    line += f"  guard={guard_state}"
            else:
                line = format_line(
                    cpu_rpm,
                    gpu_rpm,
                    cpu_temp,
                    guard_state,
                    raw if args.raw else None,
                )

            ts_iso = datetime.now().isoformat(timespec="seconds")
            if args.log_csv:
                if args.full and fields is not None:
                    row = enrich_snapshot(fields)
                    row["ts_iso"] = ts_iso
                    row["guard"] = guard_state or ""
                    append_csv_row(args.log_csv, FULL_CSV_FIELDNAMES, row)
                else:
                    append_csv_row(
                        args.log_csv,
                        MINIMAL_CSV_FIELDNAMES,
                        {
                            "ts_iso": ts_iso,
                            "cpu_rpm": cpu_rpm,
                            "gpu_rpm": gpu_rpm,
                            "cpu_temp": cpu_temp if (args.show_temp or args.guard_temp is not None) else "",
                            "guard": guard_state or "",
                        },
                    )

            if args.json:
                if args.full and fields is not None:
                    payload = enrich_snapshot(fields)
                else:
                    payload = {"cpu_rpm": cpu_rpm, "gpu_rpm": gpu_rpm}
                    if args.show_temp or args.guard_temp is not None:
                        payload["cpu_temp"] = cpu_temp
                payload["ts_iso"] = ts_iso
                if guard_state:
                    payload["guard"] = guard_state
                print(json.dumps(payload, ensure_ascii=False), flush=True)
            elif args.line or args.once:
                print(line, flush=True)
            else:
                sys.stdout.write("\r" + line + " " * 8)
                sys.stdout.flush()

            if args.once:
                break
            time.sleep(args.interval)
    except KeyboardInterrupt:
        if not args.line and not args.once and not args.json:
            print()
    except Exception as exc:
        if not args.line and not args.once:
            print()
        raise SystemExit(
            f"Failed to read fan speed: {exc}\n"
            "If this is Access denied, run the terminal as Administrator.\n"
            "If this is Not found, run with --list and pass the shown name via --instance."
        ) from exc
    finally:
        reader = None
        pythoncom.CoUninitialize()


if __name__ == "__main__":
    main()
