import csv
import time
from dataclasses import dataclass
from pathlib import Path

import pytest

from common.config_loader import load_jig_registers
from common.error_codes import RpuErrorCode, TpuErrorCode
from drivers.eload import Eload
from drivers.jig import Jig, Position
from drivers.rpu import Rpu
from drivers.tpu import Tpu

RESULTS_DIR = Path(__file__).resolve().parent.parent / "results"

# Vac_in (AC source input) is set manually by the test operator before running
# this file - not controlled by the script.
POSITION = Position(0, 0, 40)
VELOAD = 48.0
IOUT_REF = 30.0

DEFAULT_PFC_CURRENT_MAX = 30.0
DEFAULT_PFC_VDC_BUS_MIN = 350.0
DEFAULT_PFC_VDC_BUS_MAX = 425.0

FAULT_PFC_CURRENT_MAX = 8.0
FAULT_PFC_VDC_BUS_MAX = 395.0
FAULT_PFC_VDC_BUS_MIN = 445.0
AC_VOLTAGE_CALI_INCREASE_PCT = 0.10

EXPECTED_TPU_OUT_OVERCURRENT = TpuErrorCode.PFC_IL_OVER_CURRENT
EXPECTED_RPU_OUT_OVERCURRENT = RpuErrorCode.TPU_PFC_FAILURE

EXPECTED_TPU_DC_OVERVOLTAGE = TpuErrorCode.PFC_DC_OVER_VOLTAGE
EXPECTED_RPU_DC_OVERVOLTAGE = RpuErrorCode.BLE_UART_LOSS

EXPECTED_TPU_DC_UNDERVOLTAGE = TpuErrorCode.PFC_DC_UNDER_VOLTAGE
EXPECTED_RPU_DC_UNDERVOLTAGE = 0

EXPECTED_TPU_START_TIMEOUT = TpuErrorCode.PFC_START_TIMEOUT
EXPECTED_RPU_START_TIMEOUT = 0

START_CHARGING_WAIT_S = 1
CHARGING_RAMP_WAIT_S = 20  # wait for Iout to ramp up to the 30A target before recording
STOP_CHARGING_SETTLE_WAIT_S = 3
CONFIG_SETTLE_WAIT_S = 5  # settle wait for cases that don't need charging enabled

STABILITY_SAMPLE_DURATION_S = 5.0
STABILITY_SAMPLE_INTERVAL_S = 0.2
STABILITY_FIELDS = ["PFC_Vout", "PFC_iL_peak"]


def _hex32(value):
    return f"0x{value:08X}"


def _hex16(value):
    return f"0x{value:04X}"


def _error_ok(actual, actual_shadow, expected):
    return actual == expected or actual_shadow == expected


@dataclass
class PfcBaseline:
    """Shared connections + helpers for every PROT-PFC test case in this file."""
    jig: Jig
    tpu: Tpu
    rpu: Rpu
    eload: Eload
    writer: object
    csv_file: object

    def sample_min_max(self, labels, duration=STABILITY_SAMPLE_DURATION_S, interval=STABILITY_SAMPLE_INTERVAL_S):
        """Poll `labels` every `interval` seconds for `duration` seconds and
        return {label_min, label_max} - same pattern as vac_coil_estimation.py's
        sample_stability()."""
        samples = {label: [] for label in labels}
        start = time.time()
        while time.time() - start < duration:
            for label in labels:
                samples[label].append(self.tpu.read(label))
            time.sleep(interval)

        stats = {}
        for label, vals in samples.items():
            stats[f"{label}_min"] = min(vals)
            stats[f"{label}_max"] = max(vals)
        return stats

    def log_row(self, phase, setting_label, setting_value, stability_fields=STABILITY_FIELDS, pass_fail=""):
        # The 5s sampling window IS the "record" step for this row. Only the
        # fields relevant to this test case are sampled - others stay blank.
        stats = self.sample_min_max(stability_fields)
        self.writer.writerow([
            phase, setting_label, setting_value,
            stats.get("PFC_Vout_min", ""), stats.get("PFC_Vout_max", ""),
            stats.get("PFC_iL_peak_min", ""), stats.get("PFC_iL_peak_max", ""),
            self.tpu.read("RPU_output_voltage"), self.tpu.read("RPU_output_current"),
            self.tpu.read("PFC_state"), self.tpu.read("Main_state"), self.tpu.read("Main_pre_state"),
            _hex32(self.tpu.read_error_code()), _hex32(self.tpu.read_shadow_error_code()),
            _hex16(self.tpu.read("Issue code")),
            _hex32(self.rpu.read_error_code()), _hex32(self.rpu.read_shadow_error_code()),
            pass_fail,
        ])
        self.csv_file.flush()

    def start_charging(self):
        self.rpu.set_charging(True)
        time.sleep(START_CHARGING_WAIT_S)
        self.rpu.set_runtime_current(IOUT_REF)
        time.sleep(CHARGING_RAMP_WAIT_S)

    def stop_charging_and_wait_config_writable(self):
        # PFC_Current_max (and other "Debug" dataset config) is only committed
        # by the firmware while TPU is NOT actively charging (TPU_RUN_INV is
        # excluded from the config-writable states).
        self.rpu.set_charging(False)
        time.sleep(STOP_CHARGING_SETTLE_WAIT_S)
        self.tpu.wait_until_config_writable()

    def restore_config(self, label, value):
        self.stop_charging_and_wait_config_writable()
        self.tpu.write(label, value)
        self.tpu.rewrite_config()


@pytest.fixture(scope="module")
def pfc_baseline(jig_port, tpu_port, rpu_port, eload_resource):
    jig_registers = load_jig_registers()
    jig = Jig(jig_port, jig_registers)
    tpu = Tpu(tpu_port)
    rpu = Rpu(rpu_port)
    eload = Eload(eload_resource)

    RESULTS_DIR.mkdir(exist_ok=True)
    csv_path = RESULTS_DIR / f"prot_pfc_{time.strftime('%Y%m%d_%H%M%S')}.csv"
    csv_file = open(csv_path, "w", newline="", encoding="utf-8")
    writer = csv.writer(csv_file)
    writer.writerow([
        "phase", "setting_label", "setting_value", "pfc_vout_min", "pfc_vout_max",
        "pfc_il_peak_min", "pfc_il_peak_max", "rpu_output_voltage", "rpu_output_current",
        "pfc_state", "main_state", "main_pre_state", "tpu_error_code", "tpu_shadow_error_code",
        "tpu_issue_code", "rpu_error_code", "rpu_shadow_error_code", "pass_fail",
    ])

    baseline = PfcBaseline(jig=jig, tpu=tpu, rpu=rpu, eload=eload, writer=writer, csv_file=csv_file)

    try:
        rpu.set_charging(False)
        jig.move_to_position(POSITION)

        eload.set_remote()
        eload.set_mode_cv()
        eload.set_voltage(VELOAD)
        eload.output_on()

        pfc_current_max = tpu.read("PFC_Current_max")
        pfc_vdc_bus_min = tpu.read("PFC_Vdc_bus_min")
        pfc_vdc_bus_max = tpu.read("PFC_Vdc_bus_max")
        assert pfc_current_max == pytest.approx(DEFAULT_PFC_CURRENT_MAX, abs=0.5), \
            f"PFC_Current_max not at default: {pfc_current_max}"
        assert pfc_vdc_bus_min == pytest.approx(DEFAULT_PFC_VDC_BUS_MIN, abs=1), \
            f"PFC_Vdc_bus_min not at default: {pfc_vdc_bus_min}"
        assert pfc_vdc_bus_max == pytest.approx(DEFAULT_PFC_VDC_BUS_MAX, abs=1), \
            f"PFC_Vdc_bus_max not at default: {pfc_vdc_bus_max}"

        # Enable charging, wait for Iout to ramp up to 30A, then record for 5s
        # (both PFC_Vout and PFC_iL_peak are relevant as a healthy reference).
        baseline.start_charging()
        baseline.log_row("baseline", "PFC_Imax", pfc_current_max)

        yield baseline

    finally:
        rpu.set_charging(False)
        eload.output_off()
        eload.set_local()
        csv_file.close()
        jig.close()
        tpu.close()
        rpu.close()
        eload.disconnect()


def test_pfc_out_overcurrent(pfc_baseline):
    tpu = pfc_baseline.tpu
    rpu = pfc_baseline.rpu

    try:
        pfc_baseline.stop_charging_and_wait_config_writable()
        tpu.write("PFC_Current_max", FAULT_PFC_CURRENT_MAX)
        tpu.rewrite_config()

        # This case needs charging active to push real output current above
        # the lowered threshold - wait for ramp-up same as baseline.
        pfc_baseline.start_charging()

        tpu_error = tpu.read_error_code()
        tpu_shadow_error = tpu.read_shadow_error_code()
        rpu_error = rpu.read_error_code()
        rpu_shadow_error = rpu.read_shadow_error_code()

        tpu_error_ok = _error_ok(tpu_error, tpu_shadow_error, EXPECTED_TPU_OUT_OVERCURRENT)
        rpu_error_ok = _error_ok(rpu_error, rpu_shadow_error, EXPECTED_RPU_OUT_OVERCURRENT)
        pass_fail = "PASS" if (tpu_error_ok and rpu_error_ok) else "FAIL"

        # Current-protection case - only PFC_iL_peak is relevant, PFC_Vout stays blank.
        pfc_baseline.log_row("pfc_out_overcurrent", "PFC_Imax", FAULT_PFC_CURRENT_MAX,
                              stability_fields=["PFC_iL_peak"], pass_fail=pass_fail)

        assert tpu_error_ok, (
            f"TPU error code mismatch: got {_hex32(tpu_error)}/{_hex32(tpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_TPU_OUT_OVERCURRENT)}"
        )
        assert rpu_error_ok, (
            f"RPU error code mismatch: got {_hex32(rpu_error)}/{_hex32(rpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_RPU_OUT_OVERCURRENT)}"
        )

    finally:
        pfc_baseline.restore_config("PFC_Current_max", DEFAULT_PFC_CURRENT_MAX)


def test_pfc_dc_overvoltage(pfc_baseline):
    tpu = pfc_baseline.tpu
    rpu = pfc_baseline.rpu

    try:
        pfc_baseline.stop_charging_and_wait_config_writable()
        tpu.write("PFC_Vdc_bus_max", FAULT_PFC_VDC_BUS_MAX)
        tpu.rewrite_config()

        # No charging needed - the DC bus is already regulated by PFC before
        # RPU enables charging, just settle then sample.
        time.sleep(CONFIG_SETTLE_WAIT_S)

        tpu_error = tpu.read_error_code()
        tpu_shadow_error = tpu.read_shadow_error_code()
        rpu_error = rpu.read_error_code()
        rpu_shadow_error = rpu.read_shadow_error_code()

        tpu_error_ok = _error_ok(tpu_error, tpu_shadow_error, EXPECTED_TPU_DC_OVERVOLTAGE)
        rpu_error_ok = _error_ok(rpu_error, rpu_shadow_error, EXPECTED_RPU_DC_OVERVOLTAGE)
        pass_fail = "PASS" if (tpu_error_ok and rpu_error_ok) else "FAIL"

        pfc_baseline.log_row("pfc_dc_overvoltage", "Vdc_bus_max", FAULT_PFC_VDC_BUS_MAX,
                              stability_fields=["PFC_Vout"], pass_fail=pass_fail)

        assert tpu_error_ok, (
            f"TPU error code mismatch: got {_hex32(tpu_error)}/{_hex32(tpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_TPU_DC_OVERVOLTAGE)}"
        )
        assert rpu_error_ok, (
            f"RPU error code mismatch: got {_hex32(rpu_error)}/{_hex32(rpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_RPU_DC_OVERVOLTAGE)}"
        )

    finally:
        pfc_baseline.restore_config("PFC_Vdc_bus_max", DEFAULT_PFC_VDC_BUS_MAX)


def test_pfc_dc_undervoltage(pfc_baseline):
    tpu = pfc_baseline.tpu
    rpu = pfc_baseline.rpu

    try:
        pfc_baseline.stop_charging_and_wait_config_writable()
        tpu.write("PFC_Vdc_bus_min", FAULT_PFC_VDC_BUS_MIN)
        tpu.rewrite_config()

        time.sleep(CONFIG_SETTLE_WAIT_S)

        tpu_error = tpu.read_error_code()
        tpu_shadow_error = tpu.read_shadow_error_code()
        rpu_error = rpu.read_error_code()
        rpu_shadow_error = rpu.read_shadow_error_code()

        tpu_error_ok = _error_ok(tpu_error, tpu_shadow_error, EXPECTED_TPU_DC_UNDERVOLTAGE)
        rpu_error_ok = _error_ok(rpu_error, rpu_shadow_error, EXPECTED_RPU_DC_UNDERVOLTAGE)
        pass_fail = "PASS" if (tpu_error_ok and rpu_error_ok) else "FAIL"

        pfc_baseline.log_row("pfc_dc_undervoltage", "Vdc_bus_min", FAULT_PFC_VDC_BUS_MIN,
                              stability_fields=["PFC_Vout"], pass_fail=pass_fail)

        assert tpu_error_ok, (
            f"TPU error code mismatch: got {_hex32(tpu_error)}/{_hex32(tpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_TPU_DC_UNDERVOLTAGE)}"
        )
        assert rpu_error_ok, (
            f"RPU error code mismatch: got {_hex32(rpu_error)}/{_hex32(rpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_RPU_DC_UNDERVOLTAGE)}"
        )

    finally:
        pfc_baseline.restore_config("PFC_Vdc_bus_min", DEFAULT_PFC_VDC_BUS_MIN)


def test_pfc_start_timeout(pfc_baseline):
    tpu = pfc_baseline.tpu
    rpu = pfc_baseline.rpu

    # PFC_AC_Voltage_Cali_A default varies per unit - read the current value
    # first (before the try block) so `finally` can always restore it exactly,
    # even if something below fails before the fault value is written.
    pfc_baseline.stop_charging_and_wait_config_writable()
    default_cali_a = tpu.read("PFC_AC_Voltage_Cali_A")

    try:
        fault_cali_a = default_cali_a * (1 + AC_VOLTAGE_CALI_INCREASE_PCT)
        tpu.write("PFC_AC_Voltage_Cali_A", fault_cali_a)
        tpu.rewrite_config()

        time.sleep(CONFIG_SETTLE_WAIT_S)

        tpu_error = tpu.read_error_code()
        tpu_shadow_error = tpu.read_shadow_error_code()
        rpu_error = rpu.read_error_code()
        rpu_shadow_error = rpu.read_shadow_error_code()

        tpu_error_ok = _error_ok(tpu_error, tpu_shadow_error, EXPECTED_TPU_START_TIMEOUT)
        rpu_error_ok = _error_ok(rpu_error, rpu_shadow_error, EXPECTED_RPU_START_TIMEOUT)
        pass_fail = "PASS" if (tpu_error_ok and rpu_error_ok) else "FAIL"

        pfc_baseline.log_row("pfc_start_timeout", "PFC_AC_Voltage_Cali_A", fault_cali_a,
                              stability_fields=["PFC_Vout"], pass_fail=pass_fail)

        assert tpu_error_ok, (
            f"TPU error code mismatch: got {_hex32(tpu_error)}/{_hex32(tpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_TPU_START_TIMEOUT)}"
        )
        assert rpu_error_ok, (
            f"RPU error code mismatch: got {_hex32(rpu_error)}/{_hex32(rpu_shadow_error)}, "
            f"expected {_hex32(EXPECTED_RPU_START_TIMEOUT)}"
        )

    finally:
        pfc_baseline.restore_config("PFC_AC_Voltage_Cali_A", default_cali_a)