from __future__ import annotations

import inspect
import textwrap
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, ClassVar

from pyrung.core._source import (
    _capture_source,
    _capture_with_call_arg_lines,
    _capture_with_end_line,
)
from pyrung.core.condition import (
    ConditionTerm,
)
from pyrung.core.instruction import (
    ForLoopInstruction,
    SubroutineReturnSignal,
)
from pyrung.core.rung import Rung as RungLogic
from pyrung.core.tag import Tag, TagType

from .validation import _check_function_body_strict, _check_with_body_from_frame

if TYPE_CHECKING:
    from pyrung.core.context import ScanContext
    from pyrung.core.state import SystemState

    from .validation import DialectValidator

_rung_stack: list[Rung] = []


_forloop_active = False


def _current_rung() -> Rung | None:
    """Get the current rung context (if any)."""
    return _rung_stack[-1] if _rung_stack else None


def _require_rung_context(func_name: str) -> Rung:
    """Get current rung or raise error."""
    rung = _current_rung()
    if rung is None:
        raise RuntimeError(f"{func_name}() must be called inside a Rung context")
    return rung


def _push_rung_context(ctx: Rung) -> None:
    """Push a rung context onto the active stack."""
    _rung_stack.append(ctx)


def _pop_rung_context() -> None:
    """Pop the current rung context from the active stack."""
    _rung_stack.pop()


def _new_capture_context(rung: RungLogic) -> Rung:
    """Create a lightweight Rung wrapper for temporary capture scopes."""
    ctx = Rung.__new__(Rung)
    ctx._rung = rung
    ctx._pending_required_builder = None
    return ctx


def _set_scope_end_line(
    rung: RungLogic,
    *,
    source_file: str | None,
    source_line: int | None,
    context_name: str,
) -> None:
    """Populate end-line metadata for a scope-backed rung."""
    rung.end_line = _capture_with_end_line(
        source_file,
        source_line,
        context_name=context_name,
    )


def _validate_subroutine_name(name: str) -> str:
    if '"' in name:
        raise ValueError('Subroutine name must not contain ".')
    return name


class Program:
    """Container for PLC logic (rungs and subroutines).

    Used as a context manager to capture rungs:
        with Program() as logic:
            with Rung(Button):
                out(Light)

    Also works with PLCRunner:
        runner = PLCRunner(logic)
    """

    _current: Program | None = None
    _dialect_validators: ClassVar[dict[str, DialectValidator]] = {}

    def __init__(self, *, strict: bool = True) -> None:
        self._strict = strict
        self.rungs: list[RungLogic] = []
        self.subroutines: dict[str, list[RungLogic]] = {}
        self._current_subroutine: str | None = None  # Track if we're in a subroutine
        self._pending_comment: str | None = None

    def __enter__(self) -> Program:
        if self._strict:
            frame = inspect.currentframe()
            try:
                caller = frame.f_back if frame is not None else None
                if caller is not None:
                    _check_with_body_from_frame(caller, opt_out_hint="Program(strict=False)")
            finally:
                del frame
        Program._current = self
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        Program._current = None

    def add_rung(self, rung: RungLogic) -> None:
        """Add a rung to the program or current subroutine."""
        target = (
            self.subroutines[self._current_subroutine]
            if self._current_subroutine is not None
            else self.rungs
        )
        if rung._use_prior_snapshot and not target:
            scope = (
                f"subroutine '{self._current_subroutine}'"
                if self._current_subroutine is not None
                else "program"
            )
            raise RuntimeError(
                f"Rung.continued() cannot be the first rung in a {scope}. "
                "There is no prior condition snapshot to reuse."
            )
        target.append(rung)

    def start_subroutine(self, name: str) -> None:
        """Start defining a subroutine."""
        _validate_subroutine_name(name)
        self._current_subroutine = name
        self.subroutines[name] = []

    def end_subroutine(self) -> None:
        """End subroutine definition."""
        self._current_subroutine = None

    def call_subroutine(self, name: str, state: SystemState) -> SystemState:
        """Execute a subroutine by name (legacy state-based API)."""
        # This method is kept for backwards compatibility but should not be used
        # with ScanContext-based execution. Use call_subroutine_ctx instead.
        from pyrung.core.context import ScanContext

        if name not in self.subroutines:
            raise KeyError(f"Subroutine '{name}' not defined")
        ctx = ScanContext(state)
        for rung in self.subroutines[name]:
            rung.evaluate(ctx)
        return ctx.commit(dt=0.0)

    def call_subroutine_ctx(self, name: str, ctx: ScanContext) -> None:
        """Execute a subroutine by name within a ScanContext."""
        if name not in self.subroutines:
            raise KeyError(f"Subroutine '{name}' not defined")
        saved_snapshot = ctx._condition_snapshot
        ctx._condition_snapshot = None
        try:
            for rung in self.subroutines[name]:
                rung.evaluate(ctx)
        except SubroutineReturnSignal:
            pass
        finally:
            ctx._condition_snapshot = saved_snapshot

    @classmethod
    def current(cls) -> Program | None:
        """Get the current program context (if any)."""
        return cls._current

    @classmethod
    def register_dialect(cls, name: str, validator: DialectValidator) -> None:
        """Register a portability validator callback for a dialect name."""
        existing = cls._dialect_validators.get(name)
        if existing is None:
            cls._dialect_validators[name] = validator
            return
        if existing is validator:
            return
        raise ValueError(f"Dialect {name!r} already registered to a different validator")

    @classmethod
    def registered_dialects(cls) -> tuple[str, ...]:
        """Return registered dialect names in deterministic order."""
        return tuple(sorted(cls._dialect_validators))

    def validate(self, dialect: str, *, mode: str = "warn", **kwargs: Any) -> Any:
        """Run dialect-specific portability validation for this Program."""
        validator = self._dialect_validators.get(dialect)
        if validator is None:
            available = ", ".join(self.registered_dialects()) or "<none>"
            raise KeyError(
                f"Unknown validation dialect {dialect!r}. "
                f"Available dialects: {available}. "
                f"Import the dialect package first (example: import pyrung.{dialect})."
            )
        return validator(self, mode=mode, **kwargs)

    def evaluate(self, ctx: ScanContext) -> None:
        """Evaluate all main rungs in order (not subroutines) within a ScanContext."""
        for rung in self.rungs:
            rung.evaluate(ctx)


class Rung:
    """Context manager for defining a rung.

    Example:
        with Rung(Button):
            out(Light)

        with Rung(Step == 0):
            out(Light1)
            copy(1, Step, oneshot=True)
    """

    def __init__(self, *conditions: ConditionTerm) -> None:
        source_file, source_line = _capture_source(depth=2)
        self._rung = RungLogic(*conditions, source_file=source_file, source_line=source_line)
        self._pending_required_builder: tuple[int, str] | None = None
        condition_arg_lines = _capture_with_call_arg_lines(
            source_file,
            source_line,
            context_name="Rung",
        )

        # Consume any pending comment() call.
        prog = Program.current()
        if prog is not None and prog._pending_comment is not None:
            self._rung.comment = prog._pending_comment
            prog._pending_comment = None

        # Direct Tag conditions are converted internally and would otherwise
        # have no source metadata.
        for idx, condition in enumerate(self._rung._conditions):
            if condition.source_file is None:
                condition.source_file = source_file
            if condition.source_line is None:
                if idx < len(condition_arg_lines):
                    condition.source_line = condition_arg_lines[idx]
                else:
                    condition.source_line = source_line

    def continued(self) -> Rung:
        """Mark this rung as a continuation of the previous rung's condition snapshot.

        All conditions in this rung will evaluate against the same frozen
        state as the prior rung, rather than taking a fresh snapshot. This
        models multiple independent wires on the same visual rung in Click's
        ladder editor.

        Returns:
            self, for chaining: ``with Rung(B).continued(): ...``
        """
        if self._rung.comment is not None:
            raise RuntimeError(
                "A continued() rung cannot have its own comment. "
                "Set the comment on the original rung instead."
            )
        self._rung._use_prior_snapshot = True
        return self

    def _set_pending_required_builder(self, builder: object, descriptor: str) -> None:
        pending = self._pending_required_builder
        if pending is not None:
            _, existing = pending
            raise RuntimeError(f"{existing} must be completed before starting {descriptor}.")
        self._pending_required_builder = (id(builder), descriptor)

    def _assert_pending_required_builder_owner(self, builder: object, method_name: str) -> None:
        pending = self._pending_required_builder
        if pending is None:
            raise RuntimeError(
                f"{method_name}() called on a builder that is not pending in this flow."
            )
        pending_id, descriptor = pending
        if pending_id != id(builder):
            raise RuntimeError(
                f"{descriptor} must be completed before calling {method_name}() "
                "on a different builder."
            )

    def _clear_pending_required_builder(self, builder: object) -> None:
        self._assert_pending_required_builder_owner(builder, "finalize")
        self._pending_required_builder = None

    def _assert_no_pending_required_builder(self, next_action: str) -> None:
        pending = self._pending_required_builder
        if pending is None:
            return
        _, descriptor = pending
        raise RuntimeError(f"{descriptor} must be completed before calling {next_action}().")

    def _assert_required_builders_resolved(self, scope_name: str) -> None:
        pending = self._pending_required_builder
        if pending is None:
            return
        _, descriptor = pending
        raise RuntimeError(f"{descriptor} must be completed before exiting {scope_name} flow.")

    def __enter__(self) -> Rung:
        _push_rung_context(self)
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        try:
            if exc_type is None:
                self._assert_required_builders_resolved("Rung")
        finally:
            _pop_rung_context()
        _set_scope_end_line(
            self._rung,
            source_file=self._rung.source_file,
            source_line=self._rung.source_line,
            context_name="Rung",
        )
        # Add rung to current program
        prog = Program.current()
        if prog is not None:
            prog.add_rung(self._rung)


class Subroutine:
    """Context manager for defining a subroutine.

    Subroutines are named blocks of rungs that are only executed when called.

    Example:
        with subroutine("my_sub"):
            with Rung():
                out(Light)
    """

    def __init__(self, name: str, *, strict: bool = True) -> None:
        self._name = _validate_subroutine_name(name)
        self._strict = strict

    def __enter__(self) -> Subroutine:
        prog = Program.current()
        if prog is None:
            raise RuntimeError("subroutine() must be used inside a Program context")
        prog.start_subroutine(self._name)
        return self

    def __call__(self, fn: Callable[[], None]) -> SubroutineFunc:
        """Use subroutine() as a decorator.

        Example:
            @subroutine("init")
            def init_sequence():
                with Rung():
                    out(Light)
        """
        if self._strict:
            _check_function_body_strict(
                fn,
                opt_out_hint=f'@subroutine("{self._name}", strict=False)',
                source_label=f'@subroutine("{self._name}") {getattr(fn, "__qualname__", repr(fn))}',
            )
        return SubroutineFunc(self._name, fn)

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        prog = Program.current()
        if prog is not None:
            prog.end_subroutine()


class SubroutineFunc:
    """A decorated function that represents a subroutine.

    Created by using @subroutine("name") as a decorator. When passed to call(),
    auto-registers with the current Program on first use.

    Example:
        @subroutine("init")
        def init_sequence():
            with Rung():
                out(Light)

        with Program() as logic:
            with Rung(Button):
                call(init_sequence)
    """

    def __init__(self, name: str, fn: Callable[[], None]) -> None:
        self._name = name
        self._fn = fn

    @property
    def name(self) -> str:
        """The subroutine name."""
        return self._name

    def _register(self, prog: Program) -> None:
        """Register this subroutine's rungs with a Program."""
        prog.start_subroutine(self._name)
        self._fn()
        prog.end_subroutine()


def subroutine(name: str, *, strict: bool = True) -> Subroutine:
    """Define a named subroutine.

    Subroutines are only executed when called via call().
    They are NOT executed during normal program scan.

    Example:
        with Program() as logic:
            with Rung(Button):
                call("my_sub")

            with subroutine("my_sub"):
                with Rung():
                    out(Light)
    """
    return Subroutine(name, strict=strict)


class ForLoop:
    """Context manager for a repeated instruction block within a rung."""

    def __init__(
        self,
        count: Tag | int,
        oneshot: bool = False,
        source_file: str | None = None,
        source_line: int | None = None,
    ) -> None:
        self.count = count
        self.oneshot = oneshot
        self.idx = Tag("_forloop_idx", TagType.DINT)
        self._parent_ctx: Rung | None = None
        self._capture_rung: RungLogic | None = None
        self._capture_ctx: Rung | None = None
        self._source_file = source_file
        self._source_line = source_line

    def __enter__(self) -> ForLoop:
        global _forloop_active

        if _forloop_active:
            raise RuntimeError("Nested forloop is not permitted")

        self._parent_ctx = _require_rung_context("forloop")
        self._parent_ctx._assert_no_pending_required_builder("forloop")
        _forloop_active = True

        # Capture body instructions to a temporary rung (like Branch capture).
        self._capture_rung = RungLogic(
            source_file=self._source_file,
            source_line=self._source_line,
        )
        self._capture_ctx = _new_capture_context(self._capture_rung)
        _push_rung_context(self._capture_ctx)

        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        global _forloop_active

        _pop_rung_context()
        _forloop_active = False

        if exc_type is None and self._capture_ctx is not None:
            self._capture_ctx._assert_required_builders_resolved("forloop")

        if self._parent_ctx is None or self._capture_rung is None:
            return

        _set_scope_end_line(
            self._capture_rung,
            source_file=self._source_file,
            source_line=self._source_line,
            context_name="forloop",
        )

        instruction = ForLoopInstruction(
            count=self.count,
            idx_tag=self.idx,
            instructions=self._capture_rung._instructions,
            oneshot=self.oneshot,
        )
        instruction.source_file, instruction.source_line = self._source_file, self._source_line
        self._parent_ctx._rung.add_instruction(instruction)


def forloop(count: Tag | int, oneshot: bool = False) -> ForLoop:
    """Create a repeated instruction block context.

    Example:
        with Rung(Enable):
            with forloop(10) as loop:
                copy(Source[loop.idx + 1], Dest[loop.idx + 1])
    """
    source_file, source_line = _capture_source(depth=2)
    return ForLoop(count, oneshot=oneshot, source_file=source_file, source_line=source_line)


class Branch:
    """Context manager for a parallel branch within a rung.

    A branch executes when both the parent rung conditions AND
    the branch's own conditions are true.

    Example:
        with Rung(Step == 0):
            out(Light1)
            with branch(AutoMode):  # Only executes if Step==0 AND AutoMode
                out(Light2)
                copy(1, Step, oneshot=True)
    """

    def __init__(
        self,
        *conditions: ConditionTerm,
        source_file: str | None = None,
        source_line: int | None = None,
    ) -> None:
        """Create a branch with additional conditions.

        Args:
            conditions: Conditions that must be true (in addition to parent rung)
                        for this branch's instructions to execute.
        """
        self._conditions = list(conditions)
        self._branch_rung: RungLogic | None = None
        self._parent_ctx: Rung | None = None
        self._branch_ctx: Rung | None = None
        self._source_file = source_file
        self._source_line = source_line

    def __enter__(self) -> Branch:
        if _forloop_active:
            raise RuntimeError("branch() is not permitted inside forloop()")

        # Get parent rung context
        self._parent_ctx = _current_rung()
        if self._parent_ctx is None:
            raise RuntimeError("branch() must be called inside a Rung context")
        self._parent_ctx._assert_no_pending_required_builder("branch")

        # Create a nested rung for the branch that includes BOTH parent and branch conditions
        # This ensures terminal instructions (counters, timers) see the full condition chain
        parent_conditions = self._parent_ctx._rung._conditions
        condition_arg_lines = _capture_with_call_arg_lines(
            self._source_file,
            self._source_line,
            context_name="branch",
        )
        combined_conditions = parent_conditions + self._conditions
        self._branch_rung = RungLogic(
            *combined_conditions,
            source_file=self._source_file,
            source_line=self._source_line,
        )
        self._branch_rung._branch_condition_start = len(parent_conditions)

        local_conditions = self._branch_rung._conditions[
            self._branch_rung._branch_condition_start :
        ]
        for idx, condition in enumerate(local_conditions):
            if condition.source_line is None:
                if idx < len(condition_arg_lines):
                    condition.source_line = condition_arg_lines[idx]
                else:
                    condition.source_line = self._source_line

        for condition in self._branch_rung._conditions:
            if condition.source_file is None:
                condition.source_file = self._source_file
            if condition.source_line is None:
                condition.source_line = self._source_line

        # Push a new "fake" rung context so instructions go to the branch
        self._branch_ctx = _new_capture_context(self._branch_rung)
        _push_rung_context(self._branch_ctx)

        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        # Pop our branch context
        _pop_rung_context()

        if exc_type is None and self._branch_ctx is not None:
            self._branch_ctx._assert_required_builders_resolved("branch")

        # Add the branch as a nested rung to the parent
        if self._parent_ctx is not None and self._branch_rung is not None:
            _set_scope_end_line(
                self._branch_rung,
                source_file=self._source_file,
                source_line=self._source_line,
                context_name="branch",
            )
            self._parent_ctx._rung.add_branch(self._branch_rung)


def branch(*conditions: ConditionTerm) -> Branch:
    """Create a parallel branch within a rung.

    A branch executes when both the parent rung conditions AND
    the branch's own conditions are true.

    Example:
        with Rung(Step == 0):
            out(Light1)
            with branch(AutoMode):  # Only executes if Step==0 AND AutoMode
                out(Light2)
                copy(1, Step, oneshot=True)

    Args:
        conditions: Conditions that must be true (in addition to parent rung)
                    for this branch's instructions to execute.

    Returns:
        Branch context manager.
    """
    source_file, source_line = _capture_source(depth=2)
    return Branch(*conditions, source_file=source_file, source_line=source_line)


_MAX_COMMENT_LENGTH = 1400


def comment(text: str) -> None:
    """Attach a comment to the next rung.

    The comment is consumed by the immediately following ``Rung()`` constructor.
    Must be called inside a ``Program`` context.

    Example::

        comment("UnitMode Change")
        with Rung(C_UnitModeChgRequest):
            copy(1, C_UnitModeChgRequestBool, oneshot=True)
    """
    prog = Program.current()
    if prog is None:
        raise RuntimeError("comment() must be used inside a Program context")
    if prog._pending_comment is not None:
        raise RuntimeError("comment() already set — missing a Rung after the previous comment()?")
    text = textwrap.dedent(text).strip()
    if len(text) > _MAX_COMMENT_LENGTH:
        raise ValueError(f"Rung comment is {len(text)} chars, max is {_MAX_COMMENT_LENGTH}.")
    prog._pending_comment = text


RungContext = Rung
