"""Bond-graph analysis hooks.

This module intentionally has no GUI or simulation-engine dependencies.  The
causality inference algorithm can grow here without coupling the document model
to OpenModelica or Qt.
"""

from __future__ import annotations
from typing import Any

from bedit.core.application_log import get_logger

log = get_logger(__name__)

def infer_causality(component: dict[str, Any], toplevel: bool = True) -> dict[str, Any]:
    """Infer power-connection causality in a serialized component tree.

    This is currently a traversal stub: it ensures every power connection has a
    causality value, while preserving causality already supplied by callers.
    Future inference rules should assign ``source``, ``target``,
    ``warn_source``, or ``warn_target`` here and return the same tree.

    The input is mutated and returned so the composer receives the inferred
    representation without needing a second document conversion.
    """

    if toplevel:
        log.info("Causality inference")

        # Reset causalities
        reset_causality(component)

    implementation = component.get("implementation", {})
    graph = implementation.get("graph") if isinstance(implementation, dict) else None
    if not isinstance(graph, dict):
        return component
    
    id_list = build_id_list(component)

    # Fixed causalities
    for block in graph.get("blocks", []):
        for port in block.get('interface', {}).get('ports', []):
            if port.get('type', '') != 'power':
                continue
            if port.get('causality', 'indifferent') == 'fixed effort out':
                propagate_from_port(block, port, 'effort out', graph, id_list)
            elif port.get('causality', 'indifferent') == 'fixed flow out':
                propagate_from_port(block, port, 'flow out', graph, id_list)

    # Preferred causalities
    for block in graph.get("blocks", []):
        for port in block.get('interface', {}).get('ports', []):
            if port.get('type', '') != 'power':
                continue
            if is_port_fully_assigned(block, port, graph):
                continue
            if port.get('causality', 'indifferent') == 'preferred effort out':
                propagate_from_port(block, port, 'effort out', graph, id_list)
            elif port.get('causality', 'indifferent') == 'preferred flow out':
                propagate_from_port(block, port, 'flow out', graph, id_list)

            

    # Soft choices
    for block in graph.get("blocks", []):
        for port in block.get('interface', {}).get('ports', []):
            if port.get('type', '') != 'power':
                continue
            if is_port_fully_assigned(block, port, graph):
                continue
            if port.get('causality', 'indifferent') == 'likes effort out':
                propagate_from_port(block, port, 'effort out', graph, id_list)
            elif port.get('causality', 'indifferent') == 'likes flow out':
                propagate_from_port(block, port, 'flow out', graph, id_list)
            elif port.get('causality', 'indifferent') in ['indifferent', 'single flow in', 'single effort in']:
                # Force an arbitrary assignment on the first unassigned bond connected to this port
                propagate_from_port(block, port, 'effort out', graph, id_list)

    # Last check
    for con in graph.get('connections', []):
        if con.get('type') == 'power' and con.get('causality', 'none') == 'none':
            raise ValueError("System under-constrained: unresolved causal loops or disconnected elements remain")

    for block in graph.get("blocks", []):
        if isinstance(block, dict):
            infer_causality(block, False)

    return component

def reset_causality(component: dict[str, Any]) -> None:
    implementation = component.get("implementation", {})
    graph = implementation.get("graph") if isinstance(implementation, dict) else None
    if not isinstance(graph, dict):
        return
    for connection in graph.get("connections", []):
        if isinstance(connection, dict) and connection.get("type") == "power":
            connection["causality"] = "none"
    for block in graph.get("blocks", []):
        if isinstance(block, dict):
            reset_causality(block)

def build_id_list(graph: dict[str, Any]) -> dict[str, Any]:
    """Index all addressable objects in a component tree by their stable ID."""

    id_list: dict[str, Any] = {}
    id_kinds: dict[str, str] = {}

    def add(item: dict[str, Any], description: str) -> None:
        item_id = item.get("id")
        if not item_id:
            raise ValueError(f"{description} has no ID")
        # Port IDs identify a port on a component definition and may therefore
        # recur in cloned component instances. Component and junction IDs are
        # document objects and must remain globally unique.
        if item_id in id_list and not (
            description == "port" and id_kinds[item_id] == "port"
        ):
            raise ValueError(f"Duplicate simulation object ID: {item_id}")
        id_list[item_id] = item
        id_kinds[item_id] = description

    def visit(component: dict[str, Any]) -> None:
        add(component, "component")
        interface = component.get("interface", {})
        for port in interface.get("ports", []):
            add(port, "port")

        implementation = component.get("implementation", {})
        if implementation.get("kind") != "graph":
            return
        nested_graph = implementation.get("graph", {})
        for junction in nested_graph.get("junctions", []):
            add(junction, "junction")
        for block in nested_graph.get("blocks", []):
            visit(block)

    visit(graph)
    return id_list

def is_port_fully_assigned(block: dict[str, Any], port: dict[str, Any], graph: dict[str, Any]) -> bool:
    block_id = block.get('id')
    port_id = port.get('id')
    for connection in graph.get('connections', []):
        if connection.get('type') != 'power':
            continue
        source = connection.get('source', {})
        target = connection.get('target', {})
        if ((source.get('block') == block_id and source.get('port') == port_id)
                or (target.get('block') == block_id and target.get('port') == port_id)):
            if connection.get('causality', 'none') == 'none':
                return False
    return True

def propagate_from_port(block: dict[str, Any], port: dict[str, Any], desired_causality: str, graph: dict[str, Any], id_list: dict[str, Any]) -> None:
    # find all connection on this specific port
    attached_connections = []
    for con in graph.get('connections', []):
        if con.get('type') == 'power':
            if con.get('source', {}).get('block') == block.get('id') and con.get('source', {}).get('port') == port.get('id'):
                attached_connections.append(con)
            elif con.get('target', {}).get('block') == block.get('id') and con.get('target', {}).get('port') == port.get('id'):
                attached_connections.append(con)
    if not attached_connections:
        # unconnected port
        return
    
    # find first unassigned connection on this port
    target_connection = None
    for con in attached_connections:
        if con.get('causality', None) == 'none':
            target_connection = con
            break
    if target_connection is None:
        # Nothing left to resolve
        return

    # Update
    if target_connection.get('source').get('block') == block.get('id'):
        target_connection['causality'] = 'target' if (desired_causality == 'effort out') else 'source'
    else:
        target_connection['causality'] = 'source' if (desired_causality == 'effort out') else 'target'

    propagate_to_neighbor_from_connection(block, target_connection, graph, id_list)

def evaluate_junction_constraints(block: dict[str, Any], port: dict[str, Any], graph: dict[str, Any], id_list: dict[str, Any]) -> None:
    # find all connection on this specific port
    attached_connections = []
    for con in graph.get('connections', []):
        if con.get('type') == 'power':
            if con.get('source', {}).get('block') == block.get('id') and con.get('source', {}).get('port') == port.get('id'):
                attached_connections.append(con)
            elif con.get('target', {}).get('block') == block.get('id') and con.get('target', {}).get('port') == port.get('id'):
                attached_connections.append(con)
    if not attached_connections:
        # unconnected port
        return
    
    efforts_in = 0
    efforts_out = 0
    unnassigned_conns = []
    # Count current states relative to the junction port
    for con in attached_connections:
        if con.get('causality', 'none') == 'none':
            unnassigned_conns.append(con)
            continue

        is_source = (con.get('source').get('block') == block.get('id'))
        if (is_source and con.get('causality', 'none')=='source') or (not is_source and con.get('causality', 'none')=='target'):
            efforts_in += 1
        else:
            efforts_out += 1

    # Single effort in
    if port.get('causality', 'indifferent') == 'single effort in':
        # ERROR CHECK: Multiple sources/blocks trying to claim effort control on a 0-junction
        if efforts_in>1:
            raise ValueError(f"Critical causality conflict: Multiple blocks are dictating effort to 0 junction: {block.get('name')}")
        
        # 1 effort is coming in so all other ports must be effort out
        if efforts_in == 1 and len(unnassigned_conns) > 0:
            for con in unnassigned_conns:
                desired = 'target' if (con.get('source').get('block')==block.get('id')) else 'source'
                con['causality'] = desired
                propagate_to_neighbor_from_connection(block, con, graph, id_list)

        # No effort is coming in yet and one left so must be effort in
        if efforts_in == 0 and len(unnassigned_conns) == 1:
            con = unnassigned_conns[0]
            desired = 'source' if (con.get('source').get('block')==block.get('id')) else 'target'
            con['causality'] = desired
            propagate_to_neighbor_from_connection(block, con, graph, id_list)

    # Single flow in
    elif port.get('causality', 'indifferent') == 'single flow in':
        # ERROR CHECK: Multiple sources/blocks trying to claim flow control on a 1-junction
        if efforts_out>1:
            raise ValueError(f"Critical causality conflict: Multiple blocks are dictating flow to 1 junction: {block.get('name')}")
        
        # 1 flow is coming in so all other ports must be flow out
        if efforts_out == 1 and len(unnassigned_conns) > 0:
            for con in unnassigned_conns:
                desired = 'source' if (con.get('source').get('block')==block.get('id')) else 'target'
                con['causality'] = desired
                propagate_to_neighbor_from_connection(block, con, graph, id_list)

        # No flow is coming in yet and one left so must be flow in
        if efforts_out == 0 and len(unnassigned_conns) == 1:
            con = unnassigned_conns[0]
            desired = 'target' if (con.get('source').get('block')==block.get('id')) else 'source'
            con['causality'] = desired
            propagate_to_neighbor_from_connection(block, con, graph, id_list)

def propagate_to_neighbor_from_connection(block: dict[str, Any], connection: dict[str, Any], graph: dict[str, Any], id_list: dict[str, Any]) -> None:
    # Get the neighbor block and port on the other side of the bond
    neighbor_id = connection.get('target').get('block') if connection.get('source').get('block') == block.get('id') else connection.get('source').get('block') 
    neighbor_port_id = connection.get('target').get('port') if connection.get('source').get('block') == block.get('id') else connection.get('source').get('port') 
    neighbor_block = id_list.get(neighbor_id, None)
    if neighbor_block is None:
        raise ValueError(f"No other side of the bond found on connection {connection.get('id')}")
    neighbor_ports = neighbor_block.get('interface').get('ports', [])
    neighbor_port = None
    for p in neighbor_ports:
        if p.get('id') == neighbor_port_id:
            neighbor_port = p
            break
    if neighbor_port is None:
        raise ValueError(f"No other side of the bond found on connection {connection.get('id')}")
    
    # Trigger evaluation on the neighbor
    if neighbor_port.get('causality', 'indifferent') in ['single effort in', 'single flow in']:
        evaluate_junction_constraints(neighbor_block, neighbor_port, graph, id_list)
    else:
        verify_component_compatibility(neighbor_block, neighbor_port, graph)

def verify_component_compatibility(block: dict[str, Any], port: dict[str, Any], graph: dict[str, Any]) -> None:
    # Find the connection we are evaluating
    conn = None
    for c in graph.get('connections', []):
        if (c.get('source').get('block') == block.get('id') and c.get('source').get('port') == port.get('id')) or (c.get('target').get('block') == block.get('id') and c.get('target').get('port') == port.get('id')):
            if c.get('causality', 'none') != 'none':
                conn = c
                break

    if conn is None:
        return

    # Figure out what causality was pushed onto this port from the outside world
    state = conn.get('causality', 'none')
    is_source = conn.get('source', {}).get('block') == block.get('id')
    effort_out = (is_source and state == 'target') or (not is_source and state == 'source')
    port_causality = port.get('causality', 'indifferent')

    if port_causality == 'fixed effort out' and not effort_out:
        raise ValueError(f"Critical source conflict: Fixed effort source '{block.get('name')}' forced into an input state.")
    if port_causality == 'fixed flow out' and effort_out:
        raise ValueError(f"Critical source conflict: Fixed flow source '{block.get('name')}' forced into an input state.")

    if port_causality == 'preferred effort out' and not effort_out:
        conn['causality'] = f'warn_{state}'
        # log.warning(f"Storage block '{block.get('name')}' was forced into derivative causality (dependent state).")
    if port_causality == 'preferred flow out' and effort_out:
        conn['causality'] = f'warn_{state}'
        # log.warning(f"Storage block '{block.get('name')}' was forced into derivative causality (dependent state).")