from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    Literal,
    Optional,
    Sequence,
    Generator,
    overload,
)
from contextlib import contextmanager
from dataclasses import field, dataclass

from kirin.ir import Block, Region, SSAValue, Statement
from kirin.source import SourceInfo

from .abc import ASTNodeType
from .frame import Frame, CallbackFn
from .stream import StmtStream
from .exception import BuildError

if TYPE_CHECKING:
    from .abc import LoweringABC


@dataclass
class State(Generic[ASTNodeType]):
    """State of the lowering process.

    This class is used to store the state of the lowering process.
    It contains the current frame, the current block, and
    the current source.
    """

    parent: LoweringABC[ASTNodeType]
    """the parent lowering transform"""
    source: SourceInfo | None = field(default=None, kw_only=True)
    "source info of the current node"
    lines: list[str] = field(default_factory=list, kw_only=True)
    "source lines of the code being lowered"
    file: str | None = field(default=None, kw_only=True)
    "file name of the current node"
    lineno_offset: int = field(default=0, kw_only=True)
    "lineno offset at the beginning of the source"
    col_offset: int = field(default=0, kw_only=True)
    "column offset at the beginning of the source"
    _current_frame: Frame | None = field(default=None, init=False, repr=False)
    "current frame being lowered"

    def __repr__(self) -> str:
        return f"lowering.State(current_frame={self._current_frame})"

    @property
    def code(self):
        """Obtain the code generated by the lowering process."""
        region = self.root_frame.curr_region
        if not region.blocks:
            raise ValueError("No block generated")

        stmt = region.blocks[0].first_stmt
        if stmt is None:
            raise ValueError("No code generated")
        return stmt

    @property
    def root_frame(self) -> Frame[ASTNodeType]:
        """root frame of the lowering process"""
        if self._current_frame is None:
            raise ValueError("current frame is None")
        root = self._current_frame
        while root.parent is not None:
            root = root.parent
        return root

    @property
    def current_frame(self) -> Frame[ASTNodeType]:
        """current frame being lowered"""
        if self._current_frame is None:
            raise ValueError("current frame is None")
        return self._current_frame

    @dataclass
    class Result:
        """A proxy object to the result of the lowering process.

        Use `.data` to access the result of the lowering process.
        Use `.expect_one()` to assert that the result is a single value.
        """

        data: tuple[SSAValue, ...]

        def expect_one(self) -> SSAValue:
            if len(self.data) == 1:
                return self.data[0]
            raise BuildError("expected a value, but got None")

    def lower(self, node: ASTNodeType):
        result = self.parent.visit(self, node)
        if isinstance(result, Statement):
            return self.Result(tuple(result._results))
        elif result is None:
            return self.Result(tuple())
        elif isinstance(result, SSAValue):
            return self.Result((result,))
        return self.Result(result)

    def get_literal(self, value) -> SSAValue:
        return self.parent.lower_literal(self, value)

    @overload
    def get_global(
        self, node: ASTNodeType, *, no_raise: Literal[True] | bool
    ) -> LoweringABC.Result | None: ...

    @overload
    def get_global(
        self, node: ASTNodeType, *, no_raise: Literal[False] = False
    ) -> LoweringABC.Result: ...

    def get_global(self, node: ASTNodeType, *, no_raise: bool = False):
        """Get the global value of a node.
        Args:
            node (ASTNodeType): the node to get the global value of.
            no_raise (bool): if True, do not raise an exception if the value is not found.
        Returns:
            `LoweringABC.Result`: a proxy object to the global value. `.data` is the
                value, and `.expect(type)` will raise an exception if the value is the expected type.
        """
        if no_raise:
            return self.parent.lower_global_no_raise(self, node)
        return self.parent.lower_global(self, node)

    def push_frame(self, frame: Frame):
        assert frame is not self._current_frame, "Frame already pushed"
        frame.parent = self._current_frame
        self._current_frame = frame
        return frame

    def pop_frame(self, finalize_next: bool = True):
        """Pop the current frame and return it.

        Args:
            finalize_next(bool): If True, append the next block of the current frame.

        Returns:
            Frame: The popped frame.
        """
        if self._current_frame is None:
            raise ValueError("No frame to pop")
        frame = self._current_frame

        if finalize_next and frame.next_block.parent is None:
            frame.push(frame.next_block)
        self._current_frame = frame.parent
        return frame

    @contextmanager
    def frame(
        self,
        stmts: Sequence[ASTNodeType] | StmtStream[ASTNodeType],
        region: Optional[Region] = None,
        entr_block: Optional[Block] = None,
        next_block: Optional[Block] = None,
        globals: dict[str, Any] | None = None,
        capture_callback: Optional[CallbackFn] = None,
        finalize_next: bool = True,
    ) -> Generator[Frame[ASTNodeType], Any, None]:
        """Context manager to push a new frame and pop it after the block."""

        if not isinstance(stmts, StmtStream):
            stmts = StmtStream(stmts)

        region = region or Region()
        region.source = self.source

        entr_block = entr_block or Block()
        region.blocks.append(entr_block)

        if self._current_frame is not None:
            globals = globals or self.current_frame.globals
        else:
            globals = globals or {}

        frame = Frame(
            state=self,
            stream=stmts,
            curr_region=region or Region(entr_block),
            entr_block=entr_block,
            curr_block=entr_block,
            next_block=next_block or Block(),
            globals=globals,
            capture_callback=capture_callback,
        )
        self.push_frame(frame)
        try:
            yield frame
        finally:
            self.pop_frame(finalize_next)

    def error_hint(
        self,
        e: Exception,
        *,
        max_lines: int = 3,
        indent: int = 0,
        show_lineno: bool = True,
    ) -> str:
        """Generate an error hint for the given error.
        Args:
            e (Exception): The error to generate a hint for. If the error object
                has a `help` attribute, it will be used as the help message at the
                location of the error.
            max_lines (int): The maximum number of lines to show in the hint.
            indent (int): The indentation level for the hint.
            show_lineno (bool): Whether to show the line number in the hint.
        Returns:
            str: The generated error hint.
        """
        if self.source is None:
            return str(e)

        return self.source.error_hint(
            self.lines,
            e,
            file=self.file,
            indent=indent,
            show_lineno=show_lineno,
            max_lines=max_lines,
            lineno_offset=self.lineno_offset,
        )
