"""Binary data parser and temperature converter for wafer data.

Mirrors the C# Form1.cs binary parsing pipeline:
  1. Read raw bytes → strip per-block overhead → 1D hex array
  2. Chunk 1D array into rows × sensors → List[List[str]] (hex values)
  3. Convert hex values → float temperatures (family-dependent)
  4. Remove trailing zero rows
  5. Save to CSV with Sensor1, Sensor2, ... headers
"""

import logging
import os
from typing import Optional

log = logging.getLogger(__name__)

# Max DUTs (sensors) per row before overhead bytes are stripped
# P wafer: 244 valid readings per 256-block (12 overhead bytes)
# X wafer: 80 valid readings per 256-block (14 overhead bytes)
MAXDUT_P = 244
MAXDUT_X = 80


def csv_column_count(family_code: str) -> int:
    """Return the number of columns to display for a family code."""
    mapping = {
        "A": 48,
        "E": 48,
        "P": 48,
        "B": 29,
        "C": 29,
        "D": 29,
        "F": 22,
        "X": 80,
    }
    return mapping.get(family_code, 0)


def _hex_to_binary(hex_str: str) -> list[int]:
    """Convert a 4-char hex string to a 16-bit binary list (MSB first)."""
    value = int(hex_str, 16)
    return [(value >> (15 - i)) & 1 for i in range(16)]


def _twos_complement_excluding_msb(bits: list[int]) -> list[int]:
    """Invert bits 1-15, add 1 (2's complement excluding MSB)."""
    bits = list(bits)  # copy
    for i in range(1, len(bits)):
        bits[i] = 1 - bits[i]
    carry = 1
    for i in range(len(bits) - 1, 0, -1):
        s = bits[i] + carry
        bits[i] = s % 2
        carry = s // 2
    return bits


def _binary_subsequence_to_int(bits: list[int], start: int, end: int) -> int:
    """Convert bits[start:end+1] to integer."""
    result = 0
    for i in range(start, end + 1):
        result = (result << 1) | bits[i]
    return result


def _binary_fraction_to_double(bits: list[int], start: int) -> float:
    """Convert bits[start:] to a fractional value (0 < frac < 1)."""
    result = 0.0
    divisor = 2.0
    for i in range(start, len(bits)):
        result += bits[i] / divisor
        divisor *= 2.0
    return result


def _convert_standard(binary_bits: list[int]) -> float:
    """Convert 16-bit binary using standard formula (B/C/D/F families).

    Bit layout:
      bit 0  : sign
      bits 1-11 : integer part (11 bits, 2^10 .. 2^0)
      bits 12-13 : fractional part (2 bits, 2^-1, 2^-2)
      bits 14-15 : unused
    """
    value = 0.0
    for i in range(1, 14):
        if binary_bits[i]:
            value += 2.0 ** (11 - i)
    if binary_bits[0]:
        value = -value
    return value


def _convert_aep(binary_bits: list[int]) -> float:
    """Convert 16-bit binary using AEP formula.

    Bit layout:
      bit 0  : sign
      bits 1-8 : integer part (8 bits)
      bits 9-15 : fractional part (7 bits)
    """
    bits = binary_bits
    if bits[0] == 1:
        bits = _twos_complement_excluding_msb(bits)
    integer_part = _binary_subsequence_to_int(bits, 1, 8)
    fractional_part = _binary_fraction_to_double(bits, 9)
    result = integer_part + fractional_part
    if binary_bits[0] == 1:
        result = -result
    return result


def _convert_hex_to_temp(hex_str: str, family_code: str) -> float:
    """Convert a singi hale 4-char hex string to a float temperature."""
    bits = _hex_to_binary(hex_str)
    if family_code in ("A", "E", "P"):
        result = _convert_aep(bits)
    elif family_code in ("B", "C", "D"):
        result = _convert_standard(bits)
        if result < -2000:
            result = 0.0
    elif family_code == "X":
        # X wafer uses AEP-style conversion
        result = _convert_aep(bits)
    else:
        # Unknown family — try standard
        result = _convert_standard(bits)
    return round(result, 2)


def parse_binary_data(data_bytes: bytes, family_code: str) -> Optional[list[list[str]]]:
    """Parse raw wafer bytes into a 2D array of hex strings.

    Strips per-block overhead bytes (12 for P, 14 for X) and chunks
    the remaining readings into rows.

    Args:
        data_bytes: Raw binary data from the wafer.
        family_code: Wafer family code ("P", "X", "A", "B", "C", "D", "F").

    Returns:
        List of rows, each row is a list of 4-char hex strings.
        Returns None on failure.
    """
    try:
        if family_code == "X":
            return _parse_x_binary(data_bytes)
        else:
            return _parse_p_binary(data_bytes)
    except Exception as exc:
        log.error("Binary parse failed: %s", exc)
        return None


def _parse_p_binary(data_bytes: bytes) -> list[list[str]]:
    """Parse P-family (and A/B/C/D/E/F) binary data.

    Each block of 256 readings has 244 valid + 12 overhead.
    Valid readings are chunked into rows of MAXDUT_P (244).
    """
    readings: list[str] = []

    # Read 2 bytes at a time (UInt16 little-endian)
    # Each 256-word block has MAXDUT_P valid readings (first N words), rest are overhead
    num_words = len(data_bytes) // 2
    for i in range(num_words):
        value = int.from_bytes(data_bytes[i * 2 : i * 2 + 2], byteorder="little")
        if i % 256 < MAXDUT_P:
            readings.append(f"{value:04X}")

    # Chunk into rows of MAXDUT_P
    result: list[list[str]] = []
    idx = 0
    while idx + MAXDUT_P <= len(readings):
        result.append(readings[idx : idx + MAXDUT_P])
        idx += MAXDUT_P

    log.info("Parsed P-family: %d rows × %d cols", len(result), MAXDUT_P)
    return result


def _parse_x_binary(data_bytes: bytes) -> list[list[str]]:
    """Parse X-family binary data.

    Each block of 256 readings has 80 valid + 14 overhead.
    Valid readings are chunked into rows of MAXDUT_X (80).
    """
    readings: list[str] = []

    num_words = len(data_bytes) // 2
    for i in range(num_words):
        value = int.from_bytes(data_bytes[i * 2 : i * 2 + 2], byteorder="little")
        if i % 256 < MAXDUT_X:
            readings.append(f"{value:04X}")

    # Chunk into rows of MAXDUT_X
    result: list[list[str]] = []
    idx = 0
    while idx + MAXDUT_X <= len(readings):
        result.append(readings[idx : idx + MAXDUT_X])
        idx += MAXDUT_X

    log.info("Parsed X-family: %d rows × %d cols", len(result), MAXDUT_X)
    return result


def convert_to_temperatures(
    hex_data: list[list[str]], family_code: str
) -> list[list[str]]:
    """Convert hex string values to temperature strings.

    Args:
        hex_data: 2D array of 4-char hex strings from parse_binary_data.
        family_code: Wafer family code.

    Returns:
        Same structure with temperature strings (e.g. "25.37").
    """
    temp_value: list[list[str]] = []
    for row in hex_data:
        temp_row: list[str] = []
        for hex_val in row:
            temp = _convert_hex_to_temp(hex_val, family_code)
            temp_row.append(str(temp))
        temp_value.append(temp_row)
    return temp_value


def remove_trailing_zeros(data: list[list[str]]) -> None:
    """Remove rows of all-zero values from the end of data (in-place).

    A value is considered zero if its float representation is < 0.01.
    """
    while data:
        last_row = data[-1]
        is_all_zero = True
        for val in last_row:
            try:
                fval = float(val)
            except ValueError:
                fval = 99.9
            if fval >= 0.01:
                is_all_zero = False
                break
        if is_all_zero:
            data.pop()
        else:
            break


def save_to_csv(
    data: list[list[str]],
    family_code: str,
    serial_number: str,
    output_dir: str,
) -> Optional[str]:
    """Save parsed temperature data to a CSV file.

    Args:
        data: 2D array of temperature strings.
        family_code: Wafer family code.
        serial_number: Wafer serial number (e.g. "P00001").
        output_dir: Directory to save the CSV file.

    Returns:
        Full file path on success, None on failure.
    """
    try:
        os.makedirs(output_dir, exist_ok=True)
        # Build filename: P00001-20260505_133045.csv
        from datetime import datetime

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{serial_number}-{timestamp}.csv"
        filepath = os.path.join(output_dir, filename)

        num_cols = csv_column_count(family_code)
        with open(filepath, "w", encoding="utf-8") as f:
            # Header row: Sensor1,Sensor2,...
            headers = [f"Sensor{j + 1}" for j in range(num_cols)]
            f.write(",".join(headers) + "\n")
            # Data rows
            for row in data:
                # Pad or trim to match header count
                padded = row[:num_cols] + ["0"] * max(0, num_cols - len(row))
                f.write(",".join(padded) + "\n")

        log.info("Saved %d rows × %d cols to %s", len(data), num_cols, filepath)
        return filepath
    except Exception as exc:
        log.error("CSV save failed: %s", exc)
        return None