from pathlib import Path

import pytest
from hypothesis import strategies as st
from lark import Lark

from cm_inst import ArithOp, FrameOp, LogicOp, MemOp, Port, RoutingOp
from tokens import CMToken, DyadToken, FrameControlToken, MonadToken, PELocalWriteToken, SMToken

GRAMMAR_PATH = Path(__file__).parent.parent / "dfasm.lark"

uint16 = st.integers(min_value=0, max_value=0xFFFF)
int16 = st.integers(min_value=-32768, max_value=32767)
shift_amount = st.integers(min_value=0, max_value=15)

arith_dyadic_ops = st.sampled_from([ArithOp.ADD, ArithOp.SUB])
arith_monadic_ops = st.sampled_from([ArithOp.INC, ArithOp.DEC])
shift_ops = st.sampled_from([ArithOp.SHL, ArithOp.SHR, ArithOp.ASR])
logic_dyadic_ops = st.sampled_from([LogicOp.AND, LogicOp.OR, LogicOp.XOR])
comparison_ops = st.sampled_from([LogicOp.EQ, LogicOp.LT, LogicOp.LTE, LogicOp.GT, LogicOp.GTE])
branch_ops = st.sampled_from([RoutingOp.BREQ, RoutingOp.BRGT, RoutingOp.BRGE])
switch_ops = st.sampled_from([RoutingOp.SWEQ, RoutingOp.SWGT, RoutingOp.SWGE])
overflow_ops = st.sampled_from([RoutingOp.BROF, RoutingOp.SWOF])
data_routing_ops = st.sampled_from([RoutingOp.SEL, RoutingOp.MRGE])


@st.composite
def dyad_token(draw, target: int = 0, offset: int | None = None, act_id: int | None = None) -> DyadToken:
    return DyadToken(
        target=target,
        offset=draw(st.integers(min_value=0, max_value=63)) if offset is None else offset,
        act_id=draw(st.integers(min_value=0, max_value=7)) if act_id is None else act_id,
        data=draw(uint16),
        port=draw(st.sampled_from(list(Port))),
    )


@st.composite
def monad_token(draw, target: int = 0, offset: int | None = None, act_id: int | None = None) -> MonadToken:
    return MonadToken(
        target=target,
        offset=draw(st.integers(min_value=0, max_value=63)) if offset is None else offset,
        act_id=draw(st.integers(min_value=0, max_value=7)) if act_id is None else act_id,
        data=draw(uint16),
        inline=False,
    )


# SM ops: implemented for T1 are READ, WRITE, CLEAR, RD_INC, RD_DEC, CMP_SW, ALLOC, FREE, EXEC
# Unimplemented: SET_PAGE, WRITE_IMM, RAW_READ, EXT
sm_implemented_ops = [
    MemOp.READ, MemOp.WRITE, MemOp.CLEAR, MemOp.RD_INC, MemOp.RD_DEC,
    MemOp.CMP_SW, MemOp.ALLOC, MemOp.FREE, MemOp.EXEC
]
sm_all_ops = st.sampled_from(sm_implemented_ops)


@st.composite
def sm_token(draw, addr=None, op=None, data=None):
    # By default, generate T1 addresses (below tier_boundary of 256)
    # Tests that specifically need T0 addresses should pass them explicitly
    _addr = draw(st.integers(min_value=0, max_value=255)) if addr is None else addr
    _op = draw(sm_all_ops) if op is None else op
    _data = draw(uint16) if data is None else data
    ret = CMToken(target=0, offset=0, act_id=0, data=0)
    return SMToken(
        target=0,
        addr=_addr,
        op=_op,
        flags=None,
        data=_data,
        ret=ret,
    )


@st.composite
def sm_return_route(draw, target=0):
    return CMToken(
        target=target,
        offset=draw(st.integers(min_value=0, max_value=63)),
        act_id=draw(st.integers(min_value=0, max_value=7)),
        data=0,
    )


@st.composite
def pe_local_write_token(draw, target: int = 0, act_id: int | None = None) -> PELocalWriteToken:
    return PELocalWriteToken(
        target=target,
        act_id=draw(st.integers(min_value=0, max_value=7)) if act_id is None else act_id,
        region=draw(st.integers(min_value=0, max_value=1)),
        slot=draw(st.integers(min_value=0, max_value=63)),
        data=draw(uint16),
        is_dest=draw(st.booleans()),
    )


@st.composite
def frame_control_token(draw, target: int = 0, act_id: int | None = None) -> FrameControlToken:
    return FrameControlToken(
        target=target,
        act_id=draw(st.integers(min_value=0, max_value=7)) if act_id is None else act_id,
        op=draw(st.sampled_from(list(FrameOp))),
        payload=draw(uint16),
    )


@pytest.fixture(scope="session")
def parser():
    """Get the dfasm parser."""
    return Lark(
        GRAMMAR_PATH.read_text(),
        parser="earley",
        propagate_positions=True,
    )