#!/usr/bin/env python3
"""读取正点原子逻辑分析仪导出的 CSV 和 BIN 波形数据。"""

from __future__ import annotations

import array
import bisect
import csv
import io
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

from parse_logic_csv import CsvParseError, detect_delimiter, read_text, remove_comment_lines


RATE_UNITS = {
    "hz": 1.0,
    "khz": 1_000.0,
    "mhz": 1_000_000.0,
    "ghz": 1_000_000_000.0,
}
MAX_BIN_TRANSITIONS = 2_000_000
MAX_MARKER_COUNT = 10


@dataclass
class WaveformChannel:
    """保存一个数字通道的初始电平和跳变点。"""

    index: int
    name: str
    times: list[float]
    levels: list[int]

    @property
    def transition_count(self) -> int:
        """返回真实跳变次数，不包含初始点。"""
        return max(0, len(self.times) - 1)

    def average_frequency(self) -> float | None:
        """使用相邻上升沿估算平均频率。"""
        rising_times = [
            time_value
            for time_value, level in zip(self.times[1:], self.levels[1:])
            if level == 1
        ]
        if len(rising_times) < 2:
            return None
        duration = rising_times[-1] - rising_times[0]
        if duration <= 0:
            return None
        return (len(rising_times) - 1) / duration


@dataclass
class WaveformData:
    """保存一个波形文件的公共信息。"""

    path: Path
    source_type: str
    sample_rate: float | None
    sample_count: int | None
    record_count: int
    start_time: float
    end_time: float
    channels: list[WaveformChannel]

    @property
    def duration(self) -> float:
        """返回当前文件的可见时间长度。"""
        return max(0.0, self.end_time - self.start_time)


@dataclass
class FrequencyCurve:
    """保存由同方向脉冲边沿计算出的频率曲线。"""

    channel_index: int
    channel_name: str
    edge_name: str
    times: list[float]
    frequencies: list[float]

    @property
    def start_time(self) -> float:
        """返回曲线第一个有效脉冲边沿时间。"""
        return self.times[0]

    @property
    def end_time(self) -> float:
        """返回曲线最后一个有效脉冲边沿时间。"""
        return self.times[-1]

    def frequency_at(self, time_value: float) -> tuple[float, float]:
        """返回指定时间之后最近一个完整周期的时间和频率。"""
        index = bisect.bisect_left(self.times, time_value)
        if index >= len(self.times):
            index = len(self.times) - 1
        return self.times[index], self.frequencies[index]


def build_frequency_curve(
    channel: WaveformChannel, edge_mode: str = "auto"
) -> FrequencyCurve:
    """使用相邻同方向边沿的周期计算瞬时脉冲频率。"""
    if edge_mode not in ("auto", "rising", "falling"):
        raise CsvParseError(f"不支持的测量边沿：{edge_mode}")
    if len(channel.times) < 3:
        raise CsvParseError(f"{channel.name} 的脉冲边沿不足，无法计算频率")

    # times[0] 只是文件中的初始电平状态，不能当成真实边沿。
    transitions = list(zip(channel.times[1:], channel.levels[1:]))
    if edge_mode == "auto":
        selected_level = transitions[0][1]
    else:
        selected_level = 1 if edge_mode == "rising" else 0
    edge_name = "上升沿" if selected_level == 1 else "下降沿"
    edge_times = [time_value for time_value, level in transitions if level == selected_level]

    # 自动边沿不足时再尝试另一种边沿，避免文件从脉冲中间开始造成误判。
    if len(edge_times) < 2 and edge_mode == "auto":
        selected_level = 1 - selected_level
        edge_name = "上升沿" if selected_level == 1 else "下降沿"
        edge_times = [
            time_value for time_value, level in transitions if level == selected_level
        ]
    if len(edge_times) < 2:
        raise CsvParseError(f"{channel.name} 的{edge_name}不足，无法计算频率")

    frequencies: list[float] = []
    valid_edge_times: list[float] = [edge_times[0]]
    for index in range(1, len(edge_times)):
        period = edge_times[index] - edge_times[index - 1]
        if period <= 0:
            continue
        frequencies.append(1.0 / period)
        valid_edge_times.append(edge_times[index])
    if not frequencies:
        raise CsvParseError(f"{channel.name} 没有有效脉冲周期")

    # 起点还没有完整周期，使用第一个完整周期的频率补齐曲线起点。
    curve_frequencies = [frequencies[0], *frequencies]
    return FrequencyCurve(
        channel_index=channel.index,
        channel_name=channel.name,
        edge_name=edge_name,
        times=valid_edge_times,
        frequencies=curve_frequencies,
    )


def build_frequency_curves(
    channels: Iterable[WaveformChannel], edge_mode: str = "auto"
) -> list[FrequencyCurve]:
    """为所有具有足够脉冲边沿的通道生成频率曲线。"""
    curves: list[FrequencyCurve] = []
    for channel in channels:
        try:
            curves.append(build_frequency_curve(channel, edge_mode))
        except CsvParseError:
            continue
    if not curves:
        raise CsvParseError("所选通道没有足够的脉冲边沿，无法生成频率曲线")
    return curves


def parse_marker_times_ms(text: str) -> list[float]:
    """解析最多十个毫秒关键时间，并转换为秒。"""
    if not text.strip():
        return []
    markers: set[float] = set()
    for part in re.split(r"[,，\s]+", text.strip()):
        if not part:
            continue
        try:
            milliseconds = float(part)
        except ValueError as exc:
            raise CsvParseError(f"无法识别关键时间：{part!r}") from exc
        if milliseconds < 0:
            raise CsvParseError("关键时间不能小于 0 ms")
        markers.add(milliseconds / 1_000.0)
        if len(markers) > MAX_MARKER_COUNT:
            raise CsvParseError(f"关键时间最多设置 {MAX_MARKER_COUNT} 个")
    return sorted(markers)


def parse_engineering_number(text: str) -> float:
    """解析 20 MHz、500 kHz 等带单位的数值。"""
    match = re.fullmatch(
        r"\s*([0-9]+(?:\.[0-9]+)?)\s*([kKmMgG]?[hH][zZ])?\s*", text
    )
    if not match:
        raise CsvParseError(f"无法识别采样率：{text!r}")
    number = float(match.group(1))
    unit = (match.group(2) or "Hz").lower()
    return number * RATE_UNITS[unit]


def parse_scaled_count(text: str) -> int | None:
    """解析 125.351087 M 形式的采样点数量。"""
    match = re.fullmatch(r"\s*([0-9]+(?:\.[0-9]+)?)\s*([kKmMgG])?\s*", text)
    if not match:
        return None
    scales = {"": 1.0, "k": 1_000.0, "m": 1_000_000.0, "g": 1_000_000_000.0}
    return int(float(match.group(1)) * scales[(match.group(2) or "").lower()])


def parse_metadata(text: str) -> tuple[float | None, int | None]:
    """读取 CSV 注释中的采样率和采样数量。"""
    sample_rate: float | None = None
    sample_count: int | None = None
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped.startswith(";"):
            continue
        content = stripped[1:].strip()
        key, separator, value = content.partition(":")
        if not separator:
            continue
        key = key.strip().lower()
        value = value.strip()
        if key == "sample rate":
            sample_rate = parse_engineering_number(value)
        elif key == "sample count":
            sample_count = parse_scaled_count(value)
    return sample_rate, sample_count


def find_numeric_time_column(header: list[str], rows: list[list[str]]) -> int:
    """优先选择正点原子 CSV 中的 Time(s) 数字时间列。"""
    lowered = [name.strip().lower() for name in header]
    preferred_names = ("time(s)", "time (s)", "time", "timestamp", "t", "时间")
    for name in preferred_names:
        if name in lowered:
            return lowered.index(name)

    # 没有标准名称时，选择前几行都能转换成浮点数的列。
    for column in range(len(header)):
        valid = True
        for row in rows[:20]:
            if column >= len(row):
                valid = False
                break
            try:
                float(row[column].strip().lstrip("'"))
            except ValueError:
                valid = False
                break
        if valid:
            return column
    raise CsvParseError("找不到可用的数字时间列")


def find_channel_columns(
    header: list[str], rows: list[list[str]], time_index: int
) -> list[tuple[int, int, str]]:
    """识别 Channel 0、CH1、D2 等数字通道列。"""
    channels: list[tuple[int, int, str]] = []
    used_indexes: set[int] = set()
    patterns = (
        re.compile(r"^channel\s*[_-]?\s*(\d+)$", re.IGNORECASE),
        re.compile(r"^(?:ch|d)\s*[_-]?\s*(\d+)$", re.IGNORECASE),
    )

    for column, name in enumerate(header):
        if column == time_index:
            continue
        channel_index: int | None = None
        for pattern in patterns:
            match = pattern.fullmatch(name.strip())
            if match:
                channel_index = int(match.group(1))
                break
        if channel_index is None or channel_index in used_indexes:
            continue
        channels.append((column, channel_index, name.strip()))
        used_indexes.add(channel_index)

    if channels:
        return sorted(channels, key=lambda item: item[1])

    # 兼容只有“时间,电平”两列的简单文件。
    for column, name in enumerate(header):
        if column == time_index:
            continue
        valid = True
        for row in rows[:100]:
            if column >= len(row) or row[column].strip() not in ("0", "1"):
                valid = False
                break
        if valid:
            channels.append((column, len(channels), name.strip() or f"Channel {len(channels)}"))
    if not channels:
        raise CsvParseError("CSV 中没有找到高低电平通道")
    return channels


def load_csv_waveform(path: Path, encoding: str = "auto") -> WaveformData:
    """读取 CSV 中的全部数字通道并保存跳变点。"""
    text = read_text(path, encoding)
    sample_rate, sample_count = parse_metadata(text)
    clean_text = remove_comment_lines(text)
    if not clean_text:
        raise CsvParseError("CSV 中没有有效数据")

    delimiter = detect_delimiter(clean_text)
    rows = list(csv.reader(io.StringIO(clean_text), delimiter=delimiter))
    if len(rows) < 2:
        raise CsvParseError("CSV 缺少表头或采样数据")

    header = [name.strip() for name in rows[0]]
    data_rows = [row for row in rows[1:] if row]
    time_index = find_numeric_time_column(header, data_rows)
    channel_columns = find_channel_columns(header, data_rows, time_index)
    channel_map = {
        column: WaveformChannel(index, name, [], [])
        for column, index, name in channel_columns
    }

    first_time: float | None = None
    last_time: float | None = None
    previous_levels: dict[int, int] = {}
    valid_rows = 0
    for line_number, row in enumerate(data_rows, start=2):
        needed_column = max([time_index, *channel_map.keys()])
        if len(row) <= needed_column:
            raise CsvParseError(f"CSV 第 {line_number} 行列数不足")
        try:
            time_value = float(row[time_index].strip().lstrip("'"))
        except ValueError as exc:
            raise CsvParseError(f"CSV 第 {line_number} 行时间不是数字") from exc
        if last_time is not None and time_value < last_time:
            raise CsvParseError(f"CSV 第 {line_number} 行时间顺序错误")

        if first_time is None:
            first_time = time_value
        last_time = time_value
        valid_rows += 1
        for column, channel in channel_map.items():
            level_text = row[column].strip()
            if level_text not in ("0", "1"):
                raise CsvParseError(
                    f"CSV 第 {line_number} 行的 {channel.name} 不是 0 或 1"
                )
            level = int(level_text)
            if column not in previous_levels or previous_levels[column] != level:
                channel.times.append(time_value)
                channel.levels.append(level)
                previous_levels[column] = level

    if first_time is None or last_time is None:
        raise CsvParseError("CSV 中没有有效采样点")
    return WaveformData(
        path=path,
        source_type="CSV",
        sample_rate=sample_rate,
        sample_count=sample_count,
        record_count=valid_rows,
        start_time=first_time,
        end_time=last_time,
        channels=list(channel_map.values()),
    )


def parse_channel_selection(text: str, channel_count: int) -> list[int] | None:
    """解析自动、4、0,1,4、0-7 等通道选择写法。"""
    stripped = text.strip().lower()
    if not stripped or stripped in ("自动", "auto", "all", "全部"):
        return None

    selected: set[int] = set()
    for part in re.split(r"[,，\s]+", stripped):
        if not part:
            continue
        if "-" in part:
            start_text, end_text = part.split("-", 1)
            if not start_text.isdigit() or not end_text.isdigit():
                raise CsvParseError(f"无法识别通道范围：{part}")
            start = int(start_text)
            end = int(end_text)
            if start > end:
                start, end = end, start
            selected.update(range(start, end + 1))
        elif part.isdigit():
            selected.add(int(part))
        else:
            match = re.fullmatch(r"(?:channel|ch|d)\s*(\d+)", part)
            if not match:
                raise CsvParseError(f"无法识别通道：{part}")
            selected.add(int(match.group(1)))

    invalid = [index for index in sorted(selected) if not 0 <= index < channel_count]
    if invalid:
        raise CsvParseError(f"通道超出 0~{channel_count - 1} 范围：{invalid}")
    return sorted(selected)


def _sample_values(data: bytes, sample_width: int) -> Iterable[int]:
    """把原始字节按每个采样点 1 或 2 字节解释。"""
    if sample_width == 1:
        return data
    values = array.array("H")
    values.frombytes(data)
    if sys.byteorder != "little":
        values.byteswap()
    return values


def load_bin_waveform(
    path: Path,
    sample_rate: float,
    channel_count: int,
    selected_indexes: list[int] | None = None,
) -> WaveformData:
    """按无文件头的原始采样格式读取 BIN。"""
    if sample_rate <= 0:
        raise CsvParseError("BIN 采样率必须大于 0")
    if not 1 <= channel_count <= 16:
        raise CsvParseError("BIN 通道数只支持 1~16")
    sample_width = 1 if channel_count <= 8 else 2
    data = path.read_bytes()
    if not data:
        raise CsvParseError("BIN 文件为空")
    if len(data) % sample_width != 0:
        raise CsvParseError(f"BIN 文件长度不能按 {sample_width} 字节采样点整除")

    values = _sample_values(data, sample_width)
    del data
    sample_count = len(values)  # type: ignore[arg-type]
    selected = selected_indexes or list(range(channel_count))
    selected_mask = sum(1 << index for index in selected)
    first_value = values[0]  # type: ignore[index]
    channel_map = {
        index: WaveformChannel(
            index=index,
            name=f"Channel {index}",
            times=[0.0],
            levels=[(first_value >> index) & 1],
        )
        for index in selected
    }

    previous_value = first_value
    transition_total = 0
    for sample_index in range(1, sample_count):
        value = values[sample_index]  # type: ignore[index]
        changed = (value ^ previous_value) & selected_mask
        while changed:
            lowest_bit = changed & -changed
            channel_index = lowest_bit.bit_length() - 1
            channel = channel_map[channel_index]
            channel.times.append(sample_index / sample_rate)
            channel.levels.append((value >> channel_index) & 1)
            transition_total += 1
            if transition_total > MAX_BIN_TRANSITIONS:
                raise CsvParseError(
                    "BIN 跳变点过多，请在“显示通道”中只填写需要查看的通道"
                )
            changed ^= lowest_bit
        previous_value = value

    end_time = 0.0 if sample_count <= 1 else (sample_count - 1) / sample_rate
    return WaveformData(
        path=path,
        source_type="BIN",
        sample_rate=sample_rate,
        sample_count=sample_count,
        record_count=sample_count,
        start_time=0.0,
        end_time=end_time,
        channels=list(channel_map.values()),
    )


def load_waveform(
    path: Path,
    bin_sample_rate: float,
    bin_channel_count: int,
    channel_text: str,
) -> tuple[WaveformData, list[WaveformChannel]]:
    """根据扩展名读取文件，并返回需要显示的通道。"""
    suffix = path.suffix.lower()
    if suffix == ".csv":
        data = load_csv_waveform(path)
        selected_indexes = parse_channel_selection(channel_text, len(data.channels))
        if selected_indexes is None:
            visible = [channel for channel in data.channels if channel.transition_count > 0]
            if not visible:
                visible = data.channels[:8]
        else:
            selected_set = set(selected_indexes)
            visible = [channel for channel in data.channels if channel.index in selected_set]
    elif suffix == ".bin":
        selected_indexes = parse_channel_selection(channel_text, bin_channel_count)
        data = load_bin_waveform(
            path,
            sample_rate=bin_sample_rate,
            channel_count=bin_channel_count,
            selected_indexes=selected_indexes,
        )
        if selected_indexes is None:
            visible = [channel for channel in data.channels if channel.transition_count > 0]
            if not visible:
                visible = data.channels[:8]
        else:
            visible = data.channels
    else:
        raise CsvParseError("只支持正点原子导出的 CSV 或 BIN 文件")

    if not visible:
        raise CsvParseError("所选通道没有可显示的数据")
    return data, visible