from __future__ import annotations
from dataclasses import dataclass
import re
from bedit_core.models import (
    Component, 
    BondPort, 
    GraphImplementation, 
    ValueType, 
    SignalDirection, 
    Port,
    SignalPort,
    PortID,
    Parameter,
    ParameterID,
    EquationImplementation,
    BondConnection,
    SignalConnection,
    ComponentID
)

_BEVALUE_PATTERN = re.compile(r"\$([A-Za-z0-9_-]+)\$")

class CompositionError(ValueError):
    pass

@dataclass
class CompositionResult:
    modelica: str = ""
    model_name: str = ""

def compose_modelica_model(root: Component) -> CompositionResult:
    composer = Composer(root)
    return composer.compose()

class Composer():
    def __init__(self, root: Component)-> None:
        self.model : str = ""
        self.root : Component = root
        self.indent: int = 0
        self.lines: list[str] = []

    def compose(self) -> CompositionResult:
        self.indent = 0
        self.lines = []

        self.emit_model(self.root)

        # Expand macros
        macros = self.derive_macros(self.root, {})
        for i in range(len(self.lines)):
            self.lines[i] = self.expand_bevalues(self.lines[i], macros)

        self.model = '\n'.join(self.lines)
        return CompositionResult(self.model, self.model_name_for(self.root))

    def emit_model(self, component: Component) -> None:
        indentation = '\t'*self.indent
        body_indentation = '\t'*(self.indent+1)
        model_name = self.model_name_for(component)
        macros: dict[str, str] = {}

        self.lines.append(f"{indentation}model {model_name}")

        # Add BondPort definition
        if self.indent == 0:
            self.lines.extend([
                f"{body_indentation}connector BondPort",
                f"{body_indentation}\tReal e;",
                f"{body_indentation}\tflow Real f;",
                f"{body_indentation}end BondPort;",
            ])

        # Emit subcomponents
        if isinstance(component.implementation, GraphImplementation):
            for c_id, c in component.implementation.graph.components.items():
                self.indent += 1
                self.emit_model(c)
                self.indent -= 1

        # Child models have already expanded their own name-based macros.
        lines_start = len(self.lines)

        # Interface
        for port_id, port in component.interface.ports.items():
            self.lines.append(f"{body_indentation}{self.port_declaration(port)};")
            if port.multiplicity:
                macros[f"{port.name}_N"] = f"${port_id}_N$"
                macros[f"{port.name}_S"] = f"${port_id}_S$"

        # Paramters
        for param_id, param in component.parameters.items():
            self.lines.append(f"{body_indentation}{self.param_declaration(param)};")

        # declaration block if equation mode
        if isinstance(component.implementation, EquationImplementation):
            for line in component.implementation.declarations:
                self.lines.append(f"{body_indentation}{line}")

        # Component instantiations in graph mode
        if isinstance(component.implementation, GraphImplementation):
            for c_id, c in component.implementation.graph.components.items():
                self.lines.append(f"{body_indentation}{self.model_name_for(c)} {c.name};")

        # Initial equations if equation mode
        if isinstance(component.implementation, EquationImplementation):
            if len(component.implementation.initial_equations):
                self.lines.append(f"{indentation}initial equation")
                for line in component.implementation.initial_equations:
                    self.lines.append(f"{body_indentation}{line}")

        self.lines.append(f"{indentation}equation")

        # equations block if equation mode
        if isinstance(component.implementation, EquationImplementation):
            for line in component.implementation.equations:
                self.lines.append(f"{body_indentation}{line}")

        # Connections if graph mode
        elif isinstance(component.implementation, GraphImplementation):
            port_indices: dict[PortID, int] = {}
            for c_id, c in component.implementation.graph.connections.items():
                source = component.implementation.graph.components[self.get_component_from_port(component, c.source)]
                target = component.implementation.graph.components[self.get_component_from_port(component, c.target)]
                source_port_name = self.port_reference(source.interface.ports[c.source], c.source, port_indices)
                target_port_name = self.port_reference(target.interface.ports[c.target], c.target, port_indices)
                if isinstance(c, BondConnection):
                    self.lines.append(f"{body_indentation}connect({source.name}.{source_port_name}, {target.name}.{target_port_name});")
                elif isinstance(c, SignalConnection):
                    self.lines.append(f"{body_indentation}{target.name}.{target_port_name} = {source.name}.{source_port_name};")
                else:
                    raise CompositionError(f"Unknown connection type {type(c).__name__}")

        else:
            raise CompositionError(f"Unknown implementation type {type(component).__name__}")

        # Fill in macros
        for i in range(lines_start, len(self.lines)):
            try:
                self.lines[i] = self.expand_bevalues(self.lines[i], macros)
            except CompositionError:
                # Accept unresolved macros
                pass

        self.lines.append(f"{indentation}end {model_name};")

    @staticmethod
    def model_name_for(component: Component) -> str:
        # TODO make sure model name is safe
        return f"m_{component.name}"

    @staticmethod
    def modelica_type_from_ValueType(type: ValueType) -> str:
        if type == ValueType.REAL:
            return "Real"
        elif type == ValueType.INT:
            return "Integer"
        elif type == ValueType.BOOL:
            return "Boolean"
        raise CompositionError(f"Unknown value type {type}")

    @classmethod
    def port_declaration(cls, port: Port) -> str:
        direction = "output" if port.direction == SignalDirection.OUTPUT else "input"
        name = port.name
        dimensions = cls.port_dimensions(port)
        if isinstance(port, BondPort):
            return f"BondPort {name}{dimensions}"
        elif isinstance(port, SignalPort):
            port_type = cls.modelica_type_from_ValueType(port.value_type)
            #TODO units
            return f"{direction} {port_type} {name}{dimensions}"
        raise CompositionError(f"Unknown port type {type(port).__name__}")

    @classmethod
    def fixed_dimensions(cls, p: Port | Parameter) -> str:
        dimensions = p.matrix_size
        rows = max(1, p.matrix_size[0])
        columns = max(1, p.matrix_size[1])
        if rows == columns == 1:
            return ""
        if columns == 1:
            return f"[{rows}]"
        return f"[{rows},{columns}]"

    @classmethod
    def port_dimensions(cls, port: Port) -> str:
        fixed = cls.fixed_dimensions(port)
        if not port.multiplicity:
            return fixed
        entries = [f"${port.name}_N$"]
        if fixed:
            entries.extend(fixed[1:-1].split(","))
        return f"[{','.join(entries)}]"

    @staticmethod
    def expand_bevalues(text, values: dict[str, str]) -> str:
        """Replace BEdit ``$name$`` macros and reject unresolved composer values."""
        def replace(match: re.Match[str]) -> str:
            name = match.group(1)
            if name not in values:
                raise CompositionError(f"Unknown BEdit value ${name}$")
            return values[name]
        return _BEVALUE_PATTERN.sub(replace, text)        

    @classmethod
    def param_declaration(cls, param: Parameter) -> str:
        dimensions = cls.fixed_dimensions(param)
        type = cls.modelica_type_from_ValueType(param.value_type)
        # TODO units
        return f"parameter {type} {param.name}{dimensions} = {param.value}"

    def get_component_from_port(self, component: Component, port_id: PortID) -> ComponentID:
        if isinstance(component.implementation, GraphImplementation):
            for c_id, c in component.implementation.graph.components.items():
                for p_id, p in c.interface.ports.items():
                    if p_id == port_id:
                        return c_id

        raise CompositionError(f"Could not find component of port {port_id}")

    @staticmethod
    def port_reference(port: Port, port_id: PortID, indices: dict[PortID, int],) -> str:
        """Return a port name with its one-based multiplicity index."""
        if not port.multiplicity:
            return port.name
        index = indices.get(port_id, 0) + 1
        indices[port_id] = index
        return f"{port.name}[{index}]"

    def derive_macros(self, component: Component, macros: dict[str, str]) -> dict[str, str]:
        for port_id, port in component.interface.ports.items():
            if port.multiplicity:
                macros.setdefault(f"{port_id}_N", "0")
                macros.setdefault(f"{port_id}_S", "{}")

        if isinstance(component.implementation, GraphImplementation):
            for c_id, c in component.implementation.graph.connections.items():
                source = component.implementation.graph.components[self.get_component_from_port(component, c.source)]
                target = component.implementation.graph.components[self.get_component_from_port(component, c.target)]
                if source.interface.ports[c.source].multiplicity:
                    self.record_multiplicity(macros, c.source, -1)
                if target.interface.ports[c.target].multiplicity:
                    self.record_multiplicity(macros, c.target, 1)

            for c_id, c in component.implementation.graph.components.items():
                macros = self.derive_macros(c, macros)

        return macros

    @staticmethod
    def record_multiplicity(macros: dict[str, str], port_id: PortID, sign: int) -> None:
        """Record a connection count and orientation for a multiplicity port."""
        count_key = f"{port_id}_N"
        signs_key = f"{port_id}_S"
        macros[count_key] = str(int(macros.get(count_key, "0")) + 1)
        signs = macros.get(signs_key, "{}")
        entries = [] if signs == "{}" else signs[1:-1].split(", ")
        entries.append(str(float(sign)))
        macros[signs_key] = "{" + ", ".join(entries) + "}"