import re
from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Generic, TypeVar, get_type_hints

import pydantic
from tickit.core.components.component import ComponentConfig
from tickit.core.device import Device, DeviceUpdate
from tickit.core.typedefs import SimTime


@dataclass
class Param:
    """
    Representation of the internal configuration of the Zebra.
    Sometimes is shared state between multiple blocks.
    """

    reg: int
    blocks: list[str]


@dataclass
class Mux:
    """
    Representation of the internal mapping of the Zebra.
    e.g If AND1_INP1=Mux(0x08, ...) is set to 0x09, then
    AND1_INP1 should take the same input as AND1_INP2=Mux(0x09, ...)

    """

    reg: int
    block: str | None = None


DIVS = [f"DIV{i}" for i in range(1, 5)]
PULSES = [f"PULSE{i}" for i in range(1, 5)]
GATES = [f"GATE{i}" for i in range(1, 5)]

register_types: dict[str, Param | Mux] = {
    "AND1_INV": Param(0x00, ["AND1"]),
    "AND2_INV": Param(0x01, ["AND2"]),
    "AND3_INV": Param(0x02, ["AND3"]),
    "AND4_INV": Param(0x03, ["AND4"]),
    "AND1_ENA": Param(0x04, ["AND1"]),
    "AND2_ENA": Param(0x05, ["AND2"]),
    "AND3_ENA": Param(0x06, ["AND3"]),
    "AND4_ENA": Param(0x07, ["AND4"]),
    "AND1_INP1": Mux(0x08, "AND1"),
    "AND1_INP2": Mux(0x09, "AND1"),
    "AND1_INP3": Mux(0x0A, "AND1"),
    "AND1_INP4": Mux(0x0B, "AND1"),
    "AND2_INP1": Mux(0x0C, "AND2"),
    "AND2_INP2": Mux(0x0D, "AND2"),
    "AND2_INP3": Mux(0x0E, "AND2"),
    "AND2_INP4": Mux(0x0F, "AND2"),
    "AND3_INP1": Mux(0x10, "AND3"),
    "AND3_INP2": Mux(0x11, "AND3"),
    "AND3_INP3": Mux(0x12, "AND3"),
    "AND3_INP4": Mux(0x13, "AND3"),
    "AND4_INP1": Mux(0x14, "AND4"),
    "AND4_INP2": Mux(0x15, "AND4"),
    "AND4_INP3": Mux(0x16, "AND4"),
    "AND4_INP4": Mux(0x17, "AND4"),
    "OR1_INV": Param(0x18, ["OR1"]),
    "OR2_INV": Param(0x19, ["OR2"]),
    "OR3_INV": Param(0x1A, ["OR3"]),
    "OR4_INV": Param(0x1B, ["OR4"]),
    "OR1_ENA": Param(0x1C, ["OR1"]),
    "OR2_ENA": Param(0x1D, ["OR2"]),
    "OR3_ENA": Param(0x1E, ["OR3"]),
    "OR4_ENA": Param(0x1F, ["OR4"]),
    "OR1_INP1": Mux(0x20, "OR1"),
    "OR1_INP2": Mux(0x21, "OR1"),
    "OR1_INP3": Mux(0x22, "OR1"),
    "OR1_INP4": Mux(0x23, "OR1"),
    "OR2_INP1": Mux(0x24, "OR2"),
    "OR2_INP2": Mux(0x25, "OR2"),
    "OR2_INP3": Mux(0x26, "OR2"),
    "OR2_INP4": Mux(0x27, "OR2"),
    "OR3_INP1": Mux(0x28, "OR3"),
    "OR3_INP2": Mux(0x29, "OR3"),
    "OR3_INP3": Mux(0x2A, "OR3"),
    "OR3_INP4": Mux(0x2B, "OR3"),
    "OR4_INP1": Mux(0x2C, "OR4"),
    "OR4_INP2": Mux(0x2D, "OR4"),
    "OR4_INP3": Mux(0x2E, "OR4"),
    "OR4_INP4": Mux(0x2F, "OR4"),
    "GATE1_INP1": Mux(0x30, "GATE1"),
    "GATE2_INP1": Mux(0x31, "GATE2"),
    "GATE3_INP1": Mux(0x32, "GATE3"),
    "GATE4_INP1": Mux(0x33, "GATE4"),
    "GATE1_INP2": Mux(0x34, "GATE1"),
    "GATE2_INP2": Mux(0x35, "GATE2"),
    "GATE3_INP2": Mux(0x36, "GATE3"),
    "GATE4_INP2": Mux(0x37, "GATE4"),
    "DIV1_DIVLO": Param(0x38, ["DIV1"]),
    "DIV1_DIVHI": Param(0x39, ["DIV1"]),
    "DIV2_DIVLO": Param(0x3A, ["DIV2"]),
    "DIV2_DIVHI": Param(0x3B, ["DIV2"]),
    "DIV3_DIVLO": Param(0x3C, ["DIV3"]),
    "DIV3_DIVHI": Param(0x3D, ["DIV3"]),
    "DIV4_DIVLO": Param(0x3E, ["DIV4"]),
    "DIV4_DIVHI": Param(0x3F, ["DIV4"]),
    "DIV1_INP": Mux(0x40, "DIV1"),
    "DIV2_INP": Mux(0x41, "DIV2"),
    "DIV3_INP": Mux(0x42, "DIV3"),
    "DIV4_INP": Mux(0x43, "DIV4"),
    "PULSE1_DLY": Param(0x44, ["PULSE1"]),
    "PULSE2_DLY": Param(0x45, ["PULSE2"]),
    "PULSE3_DLY": Param(0x46, ["PULSE3"]),
    "PULSE4_DLY": Param(0x47, ["PULSE4"]),
    "PULSE1_WID": Param(0x48, ["PULSE1"]),
    "PULSE2_WID": Param(0x49, ["PULSE2"]),
    "PULSE3_WID": Param(0x4A, ["PULSE3"]),
    "PULSE4_WID": Param(0x4B, ["PULSE4"]),
    "PULSE1_PRE": Param(0x4C, ["PULSE1"]),
    "PULSE2_PRE": Param(0x4D, ["PULSE2"]),
    "PULSE3_PRE": Param(0x4E, ["PULSE3"]),
    "PULSE4_PRE": Param(0x4F, ["PULSE4"]),
    "PULSE1_INP": Mux(0x50, "PULSE1"),
    "PULSE2_INP": Mux(0x51, "PULSE2"),
    "PULSE3_INP": Mux(0x52, "PULSE3"),
    "PULSE4_INP": Mux(0x53, "PULSE4"),
    "POLARITY": Param(0x54, GATES + DIVS + PULSES),
    "QUAD_DIR": Param(0x55, ["QUAD"]),
    "QUAD_STEP": Param(0x56, ["QUAD"]),
    "PC_ARM_INP": Mux(0x57, "PC"),
    "PC_GATE_INP": Mux(0x58, "PC"),
    "PC_PULSE_INP": Mux(0x59, "PC"),
    "OUT1_TTL": Mux(0x60),
    "OUT1_NIM": Mux(0x61),
    "OUT1_LVDS": Mux(0x62),
    "OUT2_TTL": Mux(0x63),
    "OUT2_NIM": Mux(0x64),
    "OUT2_LVDS": Mux(0x65),
    "OUT3_TTL": Mux(0x66),
    "OUT3_OC": Mux(0x67),
    "OUT3_LVDS": Mux(0x68),
    "OUT4_TTL": Mux(0x69),
    "OUT4_NIM": Mux(0x6A),
    "OUT4_PECL": Mux(0x6B),
    "OUT5_ENCA": Mux(0x6C),
    "OUT5_ENCB": Mux(0x6D),
    "OUT5_ENCZ": Mux(0x6E),
    "OUT5_CONN": Mux(0x6F),
    "OUT6_ENCA": Mux(0x70),
    "OUT6_ENCB": Mux(0x71),
    "OUT6_ENCZ": Mux(0x72),
    "OUT6_CONN": Mux(0x73),
    "OUT7_ENCA": Mux(0x74),
    "OUT7_ENCB": Mux(0x75),
    "OUT7_ENCZ": Mux(0x76),
    "OUT7_CONN": Mux(0x77),
    "OUT8_ENCA": Mux(0x78),
    "OUT8_ENCB": Mux(0x79),
    "OUT8_ENCZ": Mux(0x7A),
    "OUT8_CONN": Mux(0x7B),
    "DIV_FIRST": Param(0x7C, DIVS),
    "SYS_RESET": Param(0x7E, []),
    "SOFT_IN": Param(0x7F, ["SOFT"]),
    "POS1_SETLO": Param(0x80, ["POS1"]),
    "POS1_SETHI": Param(0x81, ["POS1"]),
    "POS2_SETLO": Param(0x82, ["POS2"]),
    "POS2_SETHI": Param(0x83, ["POS2"]),
    "POS3_SETLO": Param(0x84, ["POS3"]),
    "POS3_SETHI": Param(0x85, ["POS3"]),
    "POS4_SETLO": Param(0x86, ["POS4"]),
    "POS4_SETHI": Param(0x87, ["POS4"]),
    "PC_ENC": Param(0x88, ["PC"]),
    "PC_TSPRE": Param(0x89, ["PC"]),
    "PC_ARM_SEL": Param(0x8A, ["PC"]),
    "PC_ARM": Param(0x8B, ["PC"]),
    "PC_DISARM": Param(0x8C, ["PC"]),
    "PC_GATE_SEL": Param(0x8D, ["PC"]),
    "PC_GATE_STARTLO": Param(0x8E, ["PC"]),
    "PC_GATE_STARTHI": Param(0x8F, ["PC"]),
    "PC_GATE_WIDLO": Param(0x90, ["PC"]),
    "PC_GATE_WIDHI": Param(0x91, ["PC"]),
    "PC_GATE_NGATELO": Param(0x92, ["PC"]),
    "PC_GATE_NGATEHI": Param(0x93, ["PC"]),
    "PC_GATE_STEPLO": Param(0x94, ["PC"]),
    "PC_GATE_STEPHI": Param(0x95, ["PC"]),
    "PC_PULSE_SEL": Param(0x96, ["PC"]),
    "PC_PULSE_STARTLO": Param(0x97, ["PC"]),
    "PC_PULSE_STARTHI": Param(0x98, ["PC"]),
    "PC_PULSE_WIDLO": Param(0x99, ["PC"]),
    "PC_PULSE_WIDHI": Param(0x9A, ["PC"]),
    "PC_PULSE_STEPLO": Param(0x9B, ["PC"]),
    "PC_PULSE_STEPHI": Param(0x9C, ["PC"]),
    "PC_PULSE_MAXLO": Param(0x9D, ["PC"]),
    "PC_PULSE_MAXHI": Param(0x9E, ["PC"]),
    "PC_BIT_CAP": Param(0x9F, ["PC"]),
    "PC_DIR": Param(0xA0, ["PC"]),
    "PC_PULSE_DLYLO": Param(0xA1, ["PC"]),
    "PC_PULSE_DLYHI": Param(0xA2, ["PC"]),
    "SYS_VER": Param(0xF0, []),
    "SYS_STATERR": Param(0xF1, []),
    "SYS_STAT1LO": Param(0xF2, []),
    "SYS_STAT1HI": Param(0xF3, []),
    "SYS_STAT2LO": Param(0xF4, []),
    "SYS_STAT2HI": Param(0xF5, []),
    "PC_NUM_CAPLO": Param(0xF6, []),
    "PC_NUM_CAPHI": Param(0xF7, []),
}

register_names = {reg.reg: name for name, reg in register_types.items()}
param_types = {name: t for name, t in register_types.items() if isinstance(t, Param)}
mux_types = {name: t for name, t in register_types.items() if isinstance(t, Mux)}

Inputs = TypeVar("Inputs")
Outputs = TypeVar("Outputs")


@dataclass
class Block(Device, ABC, Generic[Inputs, Outputs]):
    name: str
    """
    Blocks take 20ns to propagate a signal across, so we cache the next output value
    and emit it 20ns later.
    """
    previous_outputs: Outputs
    next_outputs: Outputs | None = None
    params: dict[str, int] | None = None

    def update(self, time: SimTime, inputs: Inputs) -> DeviceUpdate[Outputs]:
        if self.next_outputs:
            self.previous_outputs = self.next_outputs
            self.next_outputs = None
            return DeviceUpdate(self.previous_outputs, call_at=None)
        else:
            self.next_outputs = self._get_next_outputs(inputs)
            return DeviceUpdate(self.previous_outputs, call_at=SimTime(time + 20))

    @abstractmethod
    def _get_next_outputs(self, inputs: Inputs) -> Outputs: ...

    @property
    def num(self):
        match = re.search(r"\d*$", self.name)
        assert match, f"No trailing number in {self.name}"
        return int(match.group())

    def read_mux(self, register: str) -> int:
        return 0

    def set_mux(self, register: str, value: int) -> int:
        return 0


@pydantic.v1.dataclasses.dataclass
class BlockConfig(ComponentConfig, ABC):
    @abstractmethod
    def __call__(self) -> Block:
        """Create a block"""
        ...


def default_filler(typed_dict_type) -> Callable[[], Any]:
    def make_default():
        return {name: typ() for name, typ in get_type_hints(typed_dict_type).items()}

    return make_default


def extract_bit(registers: dict[str, int], key: str, shift: int) -> bool:
    return bool((registers[key] >> shift) & 1)


def clear_bit(registers: dict[str, int], key: str, shift: int):
    registers[key] &= ~(1 << shift)


def set_bit(registers: dict[str, int], key: str, shift: int):
    registers[key] |= 1 << shift


def rising(old: bool, new: bool) -> bool:
    return new and not old


def in_ns(ticks: int | None) -> SimTime | None:
    return None if ticks is None else SimTime(ticks * 20)
