from __future__ import annotations

from dataclasses import dataclass
from typing import TypeAlias


@dataclass
class Star:
    def __rshift__(self, other: KindExpr) -> Arrow:
        return Arrow(self, other)

    def __str__(self) -> str:
        return "*"

    def __hash__(self) -> int:
        return 0  # All Stars are the same

@dataclass
class Nat:
    def __rshift__(self, other: KindExpr) -> Arrow:
        return Arrow(self, other)

    def __str__(self) -> str:
        return "Nat"

    def __hash__(self) -> int:
        return 0  # All Stars are the same

@dataclass
class Arrow:
    """
    Represents an arrow kind `K1 -> K2`.

    To create K1 -> K2 -> K3, pass an Arrow for result_kind.
    For now, we do not support Arrows as arguments (i.e.,
    no higher order kinds).
    """
    arg_kind: Star | Nat
    result_kind: KindExpr

    def __str__(self) -> str:
        if isinstance(self.arg_kind, Star):
            arg_kind = "*"
        else:
            arg_kind = f"({str(self.arg_kind)})"

        if isinstance(self.result_kind, Star):
            result_kind = "*"
        else:
            result_kind = f"({str(self.result_kind)})"

        return f"{arg_kind} -> {result_kind}"

    def __hash__(self) -> int:
        return hash((self.arg_kind, self.result_kind, ))

KindExpr: TypeAlias = Star | Nat | Arrow