from __future__ import annotations
from typing import Dict, Any, List, Tuple, Optional
from .utils import toposort_layers, safe_get
import logging

log = logging.getLogger("opennetwork")

LAYER_MAP = {
    "dense": "Dense",
    "conv2d": "Conv2D",
    "flatten": "Flatten",
    "dropout": "Dropout",
    "batchnorm": "BatchNormalization",
    "maxpool2d": "MaxPooling2D",
    "activation": "Activation",
    "softmax_output": "Activation",  # special-cased -> 'softmax'
    "input": "Input"
}

class BuildResult:
    def __init__(self, model=None, graph_layers: List[dict] = None, inserted: List[str] = None):
        self.model = model
        self.graph_layers = graph_layers or []
        self.inserted = inserted or []  # e.g., auto-added Flatten

def _maybe_import_tf():
    try:
        import tensorflow as tf  # noqa
        return tf
    except Exception as e:
        log.debug("TensorFlow not available: %s", e)
        return None

def _layer_from_spec(tf, spec: dict):
    tname = spec.get("type", "").lower()
    params = spec.get("params", {}) or {}

    if tname == "input":
        shape = tuple(params.get("input_shape") or [])
        return tf.keras.Input(shape=shape, name=spec.get("name") or spec.get("id"))

    if tname == "dense":
        return tf.keras.layers.Dense(
            units=int(params.get("units", 32)),
            activation=params.get("activation"),
            name=spec.get("name") or spec.get("id")
        )

    if tname == "conv2d":
        ks = params.get("kernel_size") or (3, 3)
        if isinstance(ks, list):
            ks = tuple(ks)
        strides = params.get("strides") or (1, 1)
        if isinstance(strides, list):
            strides = tuple(strides)
        padding = params.get("padding") or "valid"
        activation = params.get("activation")
        return tf.keras.layers.Conv2D(
            filters=int(params.get("filters", 32)),
            kernel_size=ks,
            strides=strides,
            padding=padding,
            activation=activation,
            name=spec.get("name") or spec.get("id")
        )

    if tname == "maxpool2d":
        ps = params.get("pool_size") or (2, 2)
        if isinstance(ps, list):
            ps = tuple(ps)
        strides = params.get("strides") or None
        if isinstance(strides, list):
            strides = tuple(strides)
        return tf.keras.layers.MaxPooling2D(pool_size=ps, strides=strides, name=spec.get("name") or spec.get("id"))

    if tname == "dropout":
        return tf.keras.layers.Dropout(rate=float(params.get("rate", params.get("dropout", 0.5))), name=spec.get("name") or spec.get("id"))

    if tname == "batchnorm":
        return tf.keras.layers.BatchNormalization(name=spec.get("name") or spec.get("id"))

    if tname == "flatten":
        return tf.keras.layers.Flatten(name=spec.get("name") or spec.get("id"))

    if tname == "activation":
        return tf.keras.layers.Activation(params.get("activation", "relu"), name=spec.get("name") or spec.get("id"))

    if tname == "softmax_output":
        return tf.keras.layers.Activation("softmax", name=spec.get("name") or spec.get("id"))

    raise ValueError(f"Unsupported layer type: {tname}")

def _graph_index(layers: List[dict]):
    by_id = {n["id"]: n for n in layers}
    parents = {n["id"]: [] for n in layers}
    children = {n["id"]: (n.get("connections") or [])[:] for n in layers}
    for n in layers:
        for c in n.get("connections") or []:
            if c in parents:
                parents[c].append(n["id"])
    return by_id, parents, children

def _needs_flatten(prev_shape: Tuple[int, ...]) -> bool:
    # If previous layer outputs rank > 2 and next is Dense, we need Flatten
    if prev_shape is None:
        return False
    return len(prev_shape) > 2

def build_from_json(data: Dict[str, Any], strict: bool = False, auto_flatten: bool = True) -> BuildResult:
    """
    Build a Keras Functional model from our JSON graph if tensorflow is available.
    If tensorflow is missing, returns a BuildResult with model=None and the sorted graph.
    """
    tf = _maybe_import_tf()
    layers = data.get("layers", [])
    ordered = toposort_layers(layers)
    inserted = []

    if tf is None:
        # Dry parse only
        return BuildResult(model=None, graph_layers=ordered, inserted=inserted)

    # Build functional API
    id_to_tensor = {}
    by_id, parents, children = _graph_index(ordered)

    # Create Input tensors first
    for n in ordered:
        if n.get("type") == "input":
            tensor = _layer_from_spec(tf, n)
            id_to_tensor[n["id"]] = tensor

    # Iterate in topological order to connect tensors
    for n in ordered:
        tname = n.get("type", "").lower()
        if tname == "input":
            continue

        inbound_ids = parents.get(n["id"], [])
        inbound_tensors = [id_to_tensor[i] for i in inbound_ids if i in id_to_tensor]

        # If no inbound tensors and not input, try to chain from last seen tensor
        x = None
        if inbound_tensors:
            if len(inbound_tensors) == 1:
                x = inbound_tensors[0]
            else:
                # If multiple inbound tensors, attempt to concatenate by default
                x = tf.keras.layers.Concatenate(name=f"{n['id']}_concat")(inbound_tensors)
        else:
            # If there is exactly one existing tensor, connect to it
            if len(id_to_tensor) == 1:
                x = list(id_to_tensor.values())[0]

        # Special handling: auto-flatten before Dense if rank>2
        if tname == "dense" and auto_flatten and x is not None:
            try:
                prev_shape = x.shape.as_list()[1:]  # exclude batch
                if _needs_flatten(tuple(prev_shape)):
                    x = tf.keras.layers.Flatten(name=f"{n['id']}_auto_flatten")(x)
                    inserted.append(f"Inserted Flatten before {n['id']}")
            except Exception:
                pass

        layer_obj = _layer_from_spec(tf, n)

        # If x is None (dangling), create dummy Input based on params if possible
        if x is None:
            if tname == "dense":
                # Fallback to a flat input of size units (best-effort)
                units = int(n.get("params", {}).get("units", 32))
                inp = tf.keras.Input(shape=(units,), name=f"{n['id']}_auto_input")
                x = inp
                inserted.append(f"Inserted Input({units}) for standalone Dense {n['id']}")
            else:
                # Generic fallback input (e.g., 28x28x1) if we can guess from any upstream input
                guessed = (28, 28, 1)
                inp = tf.keras.Input(shape=guessed, name=f"{n['id']}_auto_input")
                x = inp
                inserted.append(f"Inserted default Input{guessed} for {n['id']}")

        y = layer_obj(x) if callable(layer_obj) else layer_obj

        id_to_tensor[n["id"]] = y

    # Determine outputs: nodes with no children
    output_ids = [nid for nid, kids in children.items() if len(kids) == 0]
    outputs = [id_to_tensor[nid] for nid in output_ids if nid in id_to_tensor]

    # Inputs are tensors that are tf.keras.Input
    inputs = [t for t in id_to_tensor.values() if getattr(getattr(t, "_keras_history", None), "layer", None).__class__.__name__ == "InputLayer"]

    if not inputs:
        # Fallback: use first tensor as input if no explicit Input layer
        first = next(iter(id_to_tensor.values()))
        inputs = [first]

    if not outputs:
        # Fallback: last tensor is output
        outputs = [list(id_to_tensor.values())[-1]]

    model = tf.keras.Model(inputs=inputs, outputs=outputs, name=data.get("name", "OpenNetworkModel"))
    return BuildResult(model=model, graph_layers=ordered, inserted=inserted)
