from __future__ import annotations

from collections.abc import Iterable, Sequence
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from enum import Enum
from threading import Lock
from typing import TYPE_CHECKING, Any, Literal, Self

from sonolus.backend.blocks import BlockData, PlayBlock
from sonolus.backend.ir import IRConst, IRExpr, IRStmt
from sonolus.backend.mode import Mode
from sonolus.backend.optimize.flow import BasicBlock, FlowEdge
from sonolus.backend.place import (
    PREALLOCATED_TEMP_PLACE_COUNT,
    Block,
    BlockPlace,
    TempBlock,
    preallocated_temp_block_places,
)
from sonolus.script.globals import _GlobalInfo, _GlobalPlaceholder
from sonolus.script.internal.value import Value

if TYPE_CHECKING:
    from sonolus.script.project import BuildConfig

_compiler_internal_ = True

context_var: ContextVar[Context | None] = ContextVar("context_var", default=None)  # type: ignore


@dataclass(frozen=True)
class DebugConfig:
    unchecked_reads: bool
    unchecked_writes: bool


_full_debug_config = DebugConfig(
    unchecked_reads=True,
    unchecked_writes=True,
)

_disabled_debug_config = DebugConfig(
    unchecked_reads=False,
    unchecked_writes=False,
)

debug_var = ContextVar("debug_var", default=_disabled_debug_config)


class RuntimeChecks(Enum):
    """Runtime error checking modes."""

    NONE = "none"
    """No runtime checks."""

    TERMINATE = "terminate"
    """Terminate on errors."""

    NOTIFY_AND_TERMINATE = "notify_and_terminate"
    """Log, debug pause, and terminate on errors."""


@dataclass
class FunctionVisitStatistics:
    total_time: int = 0
    own_time: int = 0
    call_count: int = 0


class ProjectContextState:
    rom: ReadOnlyMemory
    const_mappings: dict[Any, int]
    debug_str_mappings: dict[str, int]
    lock: Lock
    runtime_checks: RuntimeChecks
    visit_stats: dict[str, FunctionVisitStatistics]

    def __init__(
        self,
        rom: ReadOnlyMemory | None = None,
        const_mappings: dict[Any, int] | None = None,
        debug_str_mappings: dict[str, int] | None = None,
        runtime_checks: RuntimeChecks = RuntimeChecks.NONE,
    ):
        self.rom = ReadOnlyMemory() if rom is None else rom
        self.const_mappings = {} if const_mappings is None else const_mappings
        self.debug_str_mappings = {} if debug_str_mappings is None else debug_str_mappings
        self.lock = Lock()
        self.runtime_checks = runtime_checks
        self.visit_stats = {}

    @classmethod
    def from_build_config(
        cls,
        config: BuildConfig,
        rom: ReadOnlyMemory | None = None,
        const_mappings: dict[Any, int] | None = None,
        debug_str_mappings: dict[str, int] | None = None,
    ) -> Self:
        return cls(
            rom=rom,
            const_mappings=const_mappings,
            debug_str_mappings=debug_str_mappings,
            runtime_checks=config.runtime_checks,
        )


class ModeContextState:
    archetypes: dict[type, int]
    compile_time_only_archetypes: set[type]
    archetypes_by_name: dict[str, type]
    keys_by_archetype_id: Sequence[int]
    is_scored_by_archetype_id: Sequence[bool]
    archetype_mro_id_array_rom_indexes: Sequence[int] | None = None
    environment_mappings: dict[_GlobalInfo, int]
    environment_offsets: dict[Block, int]
    mode: Mode
    lock: Lock

    def __init__(self, mode: Mode, archetypes: list[type] | None = None):
        from sonolus.script.array import Array

        archetypes = [*archetypes] if archetypes is not None else []
        seen_archetypes = {*archetypes}
        compile_time_only_archetypes = set()
        for type_ in [*archetypes]:
            for entry in type_.mro():
                if getattr(entry, "_is_concrete_archetype_", False) and entry not in seen_archetypes:
                    archetypes.append(entry)
                    seen_archetypes.add(entry)
                    compile_time_only_archetypes.add(entry)
        self.archetypes = {type_: idx for idx, type_ in enumerate(archetypes)}
        self.compile_time_only_archetypes = compile_time_only_archetypes
        self.archetypes_by_name = {type_.name: type_ for type_, _ in self.archetypes.items()}  # type: ignore
        self.keys_by_archetype_id = (
            Array(*((getattr(a, "_key_", -1)) for a in archetypes)) if archetypes else Array[int, Literal[0]]()
        )
        self.is_scored_by_archetype_id = (
            Array(*((getattr(a, "_is_scored_", False)) for a in archetypes))
            if archetypes
            else Array[bool, Literal[0]]()
        )
        self.environment_mappings = {}
        self.environment_offsets = {}
        self.mode = mode
        self.lock = Lock()

    def _init_archetype_mro_info(self, rom: ReadOnlyMemory):
        from sonolus.script.array import Array
        from sonolus.script.num import Num

        with self.lock:
            if self.archetype_mro_id_array_rom_indexes is not None:
                return
            archetype_mro_id_values = []
            archetype_mro_id_offsets = []
            for type_ in self.archetypes:
                mro_ids = [self.archetypes[entry] for entry in type_.mro() if entry in self.archetypes]
                archetype_mro_id_offsets.append(len(archetype_mro_id_values))
                archetype_mro_id_values.append(len(mro_ids))
                archetype_mro_id_values.extend(mro_ids)
            archetype_mro_id_array_place = rom[tuple(archetype_mro_id_values)]

            archetype_mro_id_rom_indexes = Array[int, len(archetype_mro_id_offsets)]._with_value(
                [Num._accept_(offset + archetype_mro_id_array_place.index) for offset in archetype_mro_id_offsets]
            )
            self.archetype_mro_id_array_rom_indexes = archetype_mro_id_rom_indexes


class CallbackContextState:
    callback: str
    used_names: dict[str, int]
    debug_stack: list[str]
    no_eval: bool
    visitor_own_time: int

    def __init__(self, callback: str, no_eval: bool = False):
        self.callback = callback
        self.used_names = {}
        self.debug_stack = []
        self.no_eval = no_eval
        self.visitor_own_time = 0


class Context:
    project_state: ProjectContextState
    mode_state: ModeContextState
    callback_state: CallbackContextState
    statements: list[IRStmt]
    test: IRExpr
    outgoing: dict[float | None, Context]
    scope: Scope
    loop_variables: dict[str, Value]
    live: bool

    def __init__(
        self,
        project_state: ProjectContextState,
        mode_state: ModeContextState,
        callback_state: CallbackContextState,
        scope: Scope | None = None,
        live: bool = True,
        foldable_constants: dict[TempBlock, list[float | None]] | None = None,
    ):
        self.project_state = project_state
        self.mode_state = mode_state
        self.callback_state = callback_state
        self.statements = []
        self.test = IRConst(0)
        self.outgoing = {}
        self.scope = scope if scope is not None else Scope()
        self.loop_variables = {}
        self.live = live
        self.foldable_constants = foldable_constants if foldable_constants is not None else {}

    @property
    def rom(self) -> ReadOnlyMemory:
        return self.project_state.rom

    @property
    def blocks(self) -> type[Block]:
        return self.mode_state.mode.blocks

    @property
    def callback(self) -> str:
        return self.callback_state.callback

    @property
    def used_names(self) -> dict[str, int]:
        return self.callback_state.used_names

    @property
    def no_eval(self) -> bool:
        return self.callback_state.no_eval

    def check_readable(self, place: BlockPlace):
        if debug_config().unchecked_reads:
            return
        if not self.callback:
            return
        if isinstance(place.block, BlockData) and not self.is_readable(place):
            raise RuntimeError(f"Block {place.block} is not readable in {self.callback}")

    def check_writable(self, place: BlockPlace):
        if debug_config().unchecked_writes:
            return
        if not self.callback:
            return
        if isinstance(place.block, BlockData) and not self.is_writable(place):
            raise RuntimeError(f"Block {place.block} is not writable in {self.callback}")

    def is_readable(self, place: BlockPlace) -> bool:
        if debug_config().unchecked_reads:
            return True
        return self.callback and self.callback in self.blocks(place.block).readable

    def is_writable(self, place: BlockPlace) -> bool:
        if debug_config().unchecked_writes:
            return True
        return self.callback and self.callback in self.blocks(place.block).writable

    def add_statement(self, statement: IRStmt):
        if not self.live:
            return
        self.statements.append(statement)

    def add_statements(self, *statements: IRStmt):
        for statement in statements:
            self.add_statement(statement)

    def alloc(self, name: str | None = None, size: int = 1) -> BlockPlace:
        if size == 0:
            return BlockPlace(TempBlock(name or "e", 0), 0)
        name = name or ("v" if size == 1 else "a")
        num = self._get_alloc_number(name)
        if name == "v" and num < PREALLOCATED_TEMP_PLACE_COUNT:
            return preallocated_temp_block_places[num]
        return BlockPlace(TempBlock(f"{name}{num}", size), 0)

    def _get_alloc_number(self, name: str) -> int:
        if name not in self.used_names:
            self.used_names[name] = 0
        self.used_names[name] += 1
        return self.used_names[name]

    def save_alloc_state(self) -> dict[str, int]:
        return self.used_names.copy()

    def restore_alloc_state(self, state: dict[str, int]):
        self.callback_state.used_names = state.copy()

    def copy_with_scope(self, scope: Scope) -> Context:
        return Context(
            project_state=self.project_state,
            mode_state=self.mode_state,
            callback_state=self.callback_state,
            scope=scope,
            live=self.live,
            foldable_constants={**self.foldable_constants},  # We want this to persist even if the scope changes
        )

    def branch(self, condition: float | None):
        assert condition not in self.outgoing
        result = self.copy_with_scope(self.scope.copy())
        self.outgoing[condition] = result
        return result

    def branch_with_scope(self, condition: float | None, scope: Scope):
        assert condition not in self.outgoing
        result = self.copy_with_scope(scope)
        self.outgoing[condition] = result
        return result

    def new_disconnected(self):
        return self.copy_with_scope(self.scope.copy())

    def new_empty_disconnected(self):
        return self.copy_with_scope(Scope())

    def into_dead(self):
        """Create a new context for code that is unreachable, like after a return statement."""
        result = self.copy_with_scope(self.scope.copy())
        result.live = False
        return result

    def prepare_loop_header(self, to_merge: set[str]) -> Context:
        # to_merge is the set of bindings set anywhere in the loop
        # we need to invalidate them in the header if they're reference types
        # or merge them if they're value types
        # structure is self -> intermediate -> header -> body (continue -> header) | exit
        assert len(self.outgoing) == 0
        header = self.branch(None)
        for name in sorted(to_merge):
            binding = self.scope.get_binding(name)
            if not isinstance(binding, ValueBinding):
                continue
            value = binding.value
            type_ = type(value)
            if type_._is_value_type_():
                target_value = type_._from_place_(header.alloc(size=type_._size_()))
                with using_ctx(self):
                    target_value._set_(value)
                header.scope.set_value(name, target_value)
                header.loop_variables[name] = target_value
            else:
                header.scope.set_value(name, value)
                header.loop_variables[name] = value
        return header

    def branch_to_loop_header(self, header: Self):
        if not self.live:
            return
        assert len(self.outgoing) == 0
        self.outgoing[None] = header
        values = {}
        # First do a pass through and get every value
        for name, target_value in header.loop_variables.items():
            with using_ctx(self):
                if type(target_value)._is_value_type_():
                    value = self.scope.get_value(name)
                    # We make this call to _get_readonly_() to ensure that we're reading the value at this
                    # point in time specifically, since _get_readonly_ will make a copy if the value is
                    # e.g. a Num backed by a TempBlock which could be mutated.
                    values[name] = value._get_readonly_()
        # Then actually set them
        for name, target_value in header.loop_variables.items():
            with using_ctx(self):
                if type(target_value)._is_value_type_():
                    value = values[name]
                    value = type(target_value)._accept_(value)
                    target_value._set_(value)
                else:
                    value = self.scope.get_value(name)
                    if target_value is not value:
                        raise RuntimeError(
                            f"Variable '{name}' may have conflicting definitions between loop iterations"
                        )

    def map_constant(self, value: Any) -> int:
        with self.project_state.lock:
            const_mappings = self.project_state.const_mappings
            if value not in const_mappings:
                const_mappings[value] = len(const_mappings)
            return const_mappings[value]

    def map_debug_message(self, message: str) -> int:
        with self.project_state.lock:
            message_with_trace = "\n".join([*self.callback_state.debug_stack, message])
            debug_str_mappings = self.project_state.debug_str_mappings
            if message_with_trace not in debug_str_mappings:
                debug_str_mappings[message_with_trace] = len(debug_str_mappings) + 1
            return debug_str_mappings[message_with_trace]

    def get_global_base(self, value: _GlobalInfo | _GlobalPlaceholder) -> BlockPlace:
        with self.mode_state.lock:
            block = value.blocks.get(self.mode_state.mode)
            if block is None:
                raise RuntimeError(f"Global {value} is not available in '{self.mode_state.mode.name}' mode")
            if value not in self.mode_state.environment_mappings:
                if value.offset is None:
                    offset = self.mode_state.environment_offsets.get(block, 0)
                    self.mode_state.environment_mappings[value] = offset
                    self.mode_state.environment_offsets[block] = offset + value.size
                else:
                    self.mode_state.environment_mappings[value] = value.offset
            return BlockPlace(block, self.mode_state.environment_mappings[value])

    @classmethod
    def meet(cls, contexts: list[Context]) -> Context:
        if not contexts:
            raise RuntimeError("Cannot meet empty list of contexts")
        if not any(context.live for context in contexts):
            return contexts[0].into_dead()
        contexts = [context for context in contexts if context.live]
        assert not any(context.outgoing for context in contexts)
        assert all(len(context.outgoing) == 0 for context in contexts)
        target = contexts[0].copy_with_scope(Scope())
        Scope.apply_merge(target, contexts)
        for context in contexts:
            context.outgoing[None] = target
        return target

    def register_archetype(self, type_: type) -> int:
        with self.mode_state.lock:
            if type_ not in self.mode_state.archetypes:
                self.mode_state.archetypes[type_] = len(self.mode_state.archetypes)
            return self.mode_state.archetypes[type_]

    def get_archetype_mro_id_array(self, archetype_id: int) -> Sequence[int]:
        from sonolus.script.containers import ArrayPointer
        from sonolus.script.num import Num
        from sonolus.script.pointer import _deref

        self.mode_state._init_archetype_mro_info(self.rom)
        rom_index = self.mode_state.archetype_mro_id_array_rom_indexes[archetype_id]
        return ArrayPointer[int](_deref(self.blocks.EngineRom, rom_index, Num), self.blocks.EngineRom, rom_index + 1)


def ctx() -> Context | Any:  # Using Any to silence type checker warnings if it's None
    return context_var.get()


def set_ctx(value: Context | None):
    return context_var.set(value)


@contextmanager
def using_ctx(value: Context | None):
    token = context_var.set(value)
    try:
        yield
    finally:
        context_var.reset(token)


class ReadOnlyMemory:
    values: list[float]
    indexes: dict[tuple[float, ...], int]
    _lock: Lock

    def __init__(self):
        self.values = [
            float("nan"),
            float("inf"),
            float("-inf"),
        ]
        self.indexes = {}
        self._lock = Lock()

    def __getitem__(self, item: tuple[float, ...]) -> BlockPlace:
        with self._lock:
            if item not in self.indexes:
                index = len(self.values)
                self.indexes[item] = index
                self.values.extend(item)
            else:
                index = self.indexes[item]
            return BlockPlace(self.block, index)

    @property
    def block(self) -> Block:
        if ctx():
            return ctx().blocks.EngineRom
        else:
            return PlayBlock.EngineRom

    def get_value(self, index: int) -> float:
        with self._lock:
            if index < 0 or index >= len(self.values):
                raise IndexError(f"Index {index} out of bounds for ReadOnlyMemory")
            return self.values[index]


@contextmanager
def enable_debug(config: DebugConfig | None = None):
    if config is None:
        config = _full_debug_config
    token = debug_var.set(config)
    try:
        yield
    finally:
        debug_var.reset(token)


def debug_config() -> DebugConfig:
    return debug_var.get()


@dataclass
class ValueBinding:
    value: Value


@dataclass
class ConflictBinding:
    pass


@dataclass
class EmptyBinding:
    pass


Binding = ValueBinding | ConflictBinding | EmptyBinding


class Scope:
    bindings: dict[str, Binding]

    def __init__(self, bindings: dict[str, Binding] | None = None):
        self.bindings = bindings or {}

    def get_binding(self, name: str) -> Binding:
        return self.bindings.get(name, EmptyBinding())

    def set_binding(self, name: str, binding: Binding):
        self.bindings[name] = binding

    def get_value(self, name: str) -> Value | Any:
        binding = self.get_binding(name)
        match binding:
            case ValueBinding(value):
                # we don't need to call _get_() here because _set_() is never called where it could be a problem
                return value
            case ConflictBinding():
                raise RuntimeError(
                    f"Binding '{name}' has multiple conflicting definitions or may not be guaranteed to be defined"
                )
            case EmptyBinding():
                raise RuntimeError(f"Binding '{name}' is not defined")

    def set_value(self, name: str, value: Value):
        from sonolus.script.internal.impl import validate_value

        self.bindings[name] = ValueBinding(validate_value(value))

    def delete_binding(self, name: str):
        del self.bindings[name]

    def copy(self) -> Scope:
        return Scope(self.bindings.copy())

    @classmethod
    def apply_merge(cls, target: Context, incoming: list[Context]):
        if not incoming:
            return
        assert all(len(inc.outgoing) == 0 for inc in incoming)
        sources = [context.scope for context in incoming]
        keys = unique(key for source in sources for key in source.bindings)
        for key in keys:
            bindings = [source.get_binding(key) for source in sources]
            if not all(isinstance(binding, ValueBinding) for binding in bindings):
                target.scope.set_binding(key, ConflictBinding())
                continue
            values = [binding.value for binding in bindings]
            if len({id(value) for value in values}) == 1:
                target.scope.set_binding(key, ValueBinding(values[0]))
                continue
            types = {type(value) for value in values}
            if len(types) > 1:
                target.scope.set_binding(key, ConflictBinding())
                continue
            common_type: type[Value] = types.pop()
            with using_ctx(target):
                target_value = common_type._get_merge_target_(values)
            if target_value is not NotImplemented:
                for inc in incoming:
                    with using_ctx(inc):
                        target_value._set_(inc.scope.get_value(key))
                target.scope.set_value(key, target_value)
                continue
            else:
                target.scope.set_binding(key, ConflictBinding())
                continue


def iter_contexts(context: Context):
    seen = set()
    queue = [context]
    while queue:
        current = queue.pop()
        if current in seen:
            continue
        seen.add(current)
        yield current
        queue.extend(current.outgoing.values())


def context_to_cfg(context: Context) -> BasicBlock:
    blocks = {context: BasicBlock(statements=context.statements, test=context.test)}
    for current in iter_contexts(context):
        for condition, target in current.outgoing.items():
            if target not in blocks:
                blocks[target] = BasicBlock(statements=target.statements, test=target.test)
            edge = FlowEdge(src=blocks[current], dst=blocks[target], cond=condition)
            blocks[current].outgoing.add(edge)
            blocks[target].incoming.add(edge)
    result = blocks[context]
    for current in tuple(iter_contexts(context)):
        # Break cycles so memory can be cleaned without gc
        del current.outgoing
    return result


def unique[T](iterable: Iterable[T]) -> list[T]:
    result = []
    seen = set()
    for item in iterable:
        if item not in seen:
            seen.add(item)
            result.append(item)
    return result
