"""
SBMLLayout — a clean API for reading and writing the SBML Layout extension.

Design goals
------------
* Species alias nodes are a first-class concept, not an afterthought.
* All libsbml bookkeeping (glyph ids, speciesReference ids, TextGlyphs,
  SpeciesReferenceGlyphs) is handled internally.
* The resulting SBML passes libsbml validation and loads cleanly into
  SBMLDiagrams, SBMLNetwork, or any other SBML-aware tool.
* No rendering dependency — layout geometry only.

Typical usage
-------------
    import tellurium as te
    import SBMLLayout as sl

    r = te.loada('''
        J1: S1 -> S2 + A; k1*S1
        J2: S2 + A -> S3; k2*S2*A
        k1 = 0.1; k2 = 0.2
    ''')

    layout = sl.Layout(r.getSBML())

    layout.addNode('S1', x=160, y=100, w=50, h=45)
    layout.addNode('S2', x=160, y=250, w=50, h=45)
    layout.addNode('S3', x=160, y=390, w=50, h=45)
    layout.addNode('A',  x=60,  y=200, w=50, h=40)            # prime copy
    layout.addNode('A',  x=280, y=300, w=50, h=40, alias=True) # alias copy

    layout.addReaction('J1', center=(184, 205),
                       reactants=['S1'],
                       products=['S2', 'A'])    # A refers to prime (copy 0)

    layout.addReaction('J2', center=(185, 346),
                       reactants=['S2', 'A:1'], # A:1 = alias copy
                       products=['S3'])

    sbml = layout.toSBML()

Node reference syntax
---------------------
In reactant/product lists you can write:
    'S1'   — the single (or prime) copy of S1
    'A:0'  — explicitly the prime copy of A (same as 'A')
    'A:1'  — the first alias copy of A
    'A:2'  — the second alias copy of A (if it exists)

Note on bezier curves
---------------------
SBMLLayout writes structural layout (node positions, reaction connectivity,
alias wiring) into the SBML Layout extension. Bezier curve control points
are not written into the SBML — they must be set after loading into a
visualisation tool via that tool's own API (e.g. SBMLDiagrams'
setReactionBezierHandles). This is a limitation of the SBML Layout
extension's support in current Python bindings.
"""

from __future__ import annotations

import libsbml
from dataclasses import dataclass
from typing import Optional


# ---------------------------------------------------------------------------
# Internal data structures
# ---------------------------------------------------------------------------

@dataclass
class _NodeRecord:
    """Represents one visual copy of a species node."""
    species_id: str   # SBML species id (same for prime and all aliases)
    copy_index: int   # 0 = prime, 1 = first alias, etc.
    x: float
    y: float
    w: float
    h: float
    glyph_id: str = ''  # set by Layout._assign_glyph_ids()


@dataclass
class _ReactionRecord:
    """Represents one reaction with its junction position and optional bezier handles."""
    reaction_id: str
    center: Optional[tuple]   # (x, y) junction point; None = no junction (direct reaction)
    reactant_refs: list       # list of (species_id, copy_index) tuples
    product_refs: list        # list of (species_id, copy_index) tuples
    # Optional bezier handles per arm.
    # Each entry is ((bp1x,bp1y), (bp2x,bp2y)) or None for a straight arm.
    # Order: reactant arms first, then product arms, matching reactant_refs/product_refs.
    handles: Optional[list] = None
    # If True, this is a direct reaction (no junction node). For a uni-uni direct
    # bezier, handles[0] is ((BP1x,BP1y),(BP2x,BP2y)) for the single curve from
    # reactant to product. The junction circle is not drawn.
    direct: bool = False


@dataclass
class _CompartmentRecord:
    x: float
    y: float
    w: float
    h: float


# ---------------------------------------------------------------------------
# Layout class
# ---------------------------------------------------------------------------

class Layout:
    """
    Attaches or replaces the SBML Layout extension on an SBML model.

    Parameters
    ----------
    sbml_string : str
        A valid SBML Level 3 string (e.g. from ``tellurium`` or ``libsbml``).
        Any existing layout information is ignored and replaced when
        ``toSBML()`` is called.
    """

    def __init__(self, sbml_string: str):
        self._sbml_string = sbml_string

        # Validate the base SBML
        doc = libsbml.readSBMLFromString(sbml_string)
        if doc.getNumErrors(libsbml.LIBSBML_SEV_FATAL) > 0:
            raise ValueError(
                f"Invalid SBML: {doc.getError(0).getMessage()}")
        model = doc.getModel()
        if model is None:
            raise ValueError("SBML string contains no model.")

        # Cache species and reaction ids for validation
        self._species_ids = {
            model.getSpecies(i).getId()
            for i in range(model.getNumSpecies())
        }
        self._reaction_ids = {
            model.getReaction(i).getId()
            for i in range(model.getNumReactions())
        }

        self._nodes: list[_NodeRecord] = []
        self._reactions: list[_ReactionRecord] = []
        self._compartments: dict[str, _CompartmentRecord] = {}

        # Track how many copies of each species have been added
        self._copy_counts: dict[str, int] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def addNode(self, species_id: str,
                x: float, y: float,
                w: float = 60.0, h: float = 30.0,
                alias: bool = False) -> 'Layout':
        """
        Add a species node to the layout.

        Parameters
        ----------
        species_id : str
            Must match a species id in the SBML model.
        x, y : float
            Top-left corner of the node bounding box.
        w, h : float
            Width and height of the node. Defaults to 60 x 30.
        alias : bool
            If True this is an additional copy of ``species_id``.
            Copy indices are assigned in call order: the first call for a
            given ``species_id`` is copy 0 (the prime), each subsequent
            call increments the index (copy 1, copy 2, …).

        Returns
        -------
        Layout
            Returns self for method chaining.
        """
        if species_id not in self._species_ids:
            raise ValueError(
                f"Species '{species_id}' not found in SBML model. "
                f"Available: {sorted(self._species_ids)}")

        copy_index = self._copy_counts.get(species_id, 0)
        self._copy_counts[species_id] = copy_index + 1

        self._nodes.append(_NodeRecord(
            species_id=species_id,
            copy_index=copy_index,
            x=x, y=y, w=w, h=h
        ))
        return self

    def addReaction(self, reaction_id: str,
                    reactants: list,
                    products: list,
                    center: Optional[tuple] = None,
                    handles: Optional[list] = None,
                    direct: bool = False) -> 'Layout':
        """
        Add a reaction to the layout.

        Parameters
        ----------
        reaction_id : str
            Must match a reaction id in the SBML model.
        reactants : list of str
            Species references for the reactant arms. Each entry is either
            a plain species id (``'S1'``) or a species id with copy index
            (``'A:1'`` for the second copy of A).
        products : list of str
            Species references for the product arms, same syntax as reactants.
        center : (x, y) tuple, optional
            The junction point of the reaction. If None, no junction position
            is set and the visualisation tool will compute a default.
        handles : list of ((bp1x,bp1y),(bp2x,bp2y)) or None, optional
            Bezier control points, one pair per arm in the order:
            [reactant_0, reactant_1, ..., product_0, product_1, ...].
            Each entry is a pair of tuples ``((bp1x, bp1y), (bp2x, bp2y))``
            giving the two cubic bezier control points for that arm, where
            BP1 is near the species node and BP2 is near the junction.
            Use ``None`` for a specific arm to draw it as a straight line.
            If the ``handles`` parameter itself is ``None``, all arms are
            straight lines.

        Returns
        -------
        Layout
            Returns self for method chaining.

        Examples
        --------
        Bi-uni reaction S1 + S2 -> S3 with curved reactant arms::

            layout.addReaction('J1',
                reactants=['S1', 'S2'],
                products=['S3'],
                center=(250, 200),
                handles=[
                    ((210, 102), (240, 195)),  # S1 arm: BP1 near S1, BP2 near junction
                    ((210, 302), (240, 211)),  # S2 arm: BP1 near S2, BP2 near junction
                    None,                       # S3 arm: straight line
                ])
        """
        if reaction_id not in self._reaction_ids:
            raise ValueError(
                f"Reaction '{reaction_id}' not found in SBML model. "
                f"Available: {sorted(self._reaction_ids)}")

        reactant_refs = [self._parse_species_ref(r) for r in reactants]
        product_refs  = [self._parse_species_ref(p) for p in products]

        for sid, cidx in reactant_refs + product_refs:
            self._validate_node_ref(sid, cidx, reaction_id)

        self._reactions.append(_ReactionRecord(
            reaction_id=reaction_id,
            center=center,
            reactant_refs=reactant_refs,
            product_refs=product_refs,
            handles=handles,
            direct=direct,
        ))
        return self

    def addCompartment(self, compartment_id: str,
                       x: float, y: float,
                       w: float, h: float) -> 'Layout':
        """
        Add a compartment bounding box to the layout.

        Parameters
        ----------
        compartment_id : str
            Must match a compartment id in the SBML model.
        x, y, w, h : float
            Position and size of the compartment bounding box.

        Returns
        -------
        Layout
            Returns self for method chaining.
        """
        self._compartments[compartment_id] = _CompartmentRecord(x, y, w, h)
        return self

    def draw(self,
             output=None,
             scale: float = 1.0,
             style=None) -> None:
        """
        Render the network and display or save it.

        Convenience wrapper around ``Renderer.draw()``.  Calls ``toSBML()``
        internally so no separate export step is needed.

        Parameters
        ----------
        output : str, optional
            File path (``.png``, ``.pdf``, or ``.svg``).
            If None, displays inline in Jupyter / opens in default viewer.
        scale : float
            Scale factor for raster (PNG) output. Use 2.0 for high-DPI.
        style : Style, optional
            Visual style overrides. If None, uses default Style.
        """
        from .renderer import Renderer
        Renderer(self.toSBML()).draw(output=output, scale=scale, style=style)

    def toSBML(self) -> str:
        """
        Return a valid SBML Level 3 string with the layout extension embedded.

        Handles all bookkeeping internally:

        - Sets species ``name`` attributes equal to their ids so downstream
          tools can display labels without extra configuration.
        - Assigns ids to ``speciesReference`` elements (required for valid
          ``SpeciesReferenceGlyph`` entries).
        - Creates one ``SpeciesGlyph`` + one ``TextGlyph`` per node (including
          alias copies). ``TextGlyphs`` are required by SBMLDiagrams.
        - Wires each reaction arm to the correct glyph copy via
          ``speciesGlyphId``, ensuring alias nodes connect to the right
          reaction.

        Returns
        -------
        str
            SBML Level 3 string with the Layout package included.
        """
        doc = libsbml.readSBMLFromString(self._sbml_string)
        model = doc.getModel()

        # Set species name = id so downstream tools display labels
        for i in range(model.getNumSpecies()):
            sp = model.getSpecies(i)
            if not sp.getName():
                sp.setName(sp.getId())

        # Assign ids to speciesReferences so SpeciesReferenceGlyphs are valid
        for i in range(model.getNumReactions()):
            rxn = model.getReaction(i)
            rid = rxn.getId()
            for j in range(rxn.getNumReactants()):
                sr = rxn.getReactant(j)
                if not sr.getId():
                    sr.setId(f'sr_{rid}_r{j}')
            for j in range(rxn.getNumProducts()):
                sr = rxn.getProduct(j)
                if not sr.getId():
                    sr.setId(f'sr_{rid}_p{j}')

        # Enable layout package and create a fresh layout
        doc.enablePackage(
            'http://www.sbml.org/sbml/level3/version1/layout/version1',
            'layout', True)
        doc.setPackageRequired('layout', False)
        mplugin = model.getPlugin('layout')
        while mplugin.getNumLayouts() > 0:
            mplugin.removeLayout(0)
        layout = mplugin.createLayout()
        layout.setId('layout1')

        self._assign_glyph_ids()
        self._build_compartment_glyphs(layout)
        self._build_species_glyphs(layout)
        self._build_reaction_glyphs(model, layout)

        return libsbml.writeSBMLToString(doc)

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _parse_species_ref(self, ref: str) -> tuple[str, int]:
        """Parse 'A:1' -> ('A', 1),  'S1' -> ('S1', 0)."""
        if ':' in ref:
            parts = ref.split(':', 1)
            return parts[0], int(parts[1])
        return ref, 0

    def _validate_node_ref(self, species_id: str, copy_index: int,
                           reaction_id: str):
        """Raise ValueError if no node record matches species_id + copy_index."""
        for node in self._nodes:
            if node.species_id == species_id and node.copy_index == copy_index:
                return
        copies = [n.copy_index for n in self._nodes
                  if n.species_id == species_id]
        if not copies:
            raise ValueError(
                f"Reaction '{reaction_id}': species '{species_id}' has no "
                f"node added. Call addNode('{species_id}', ...) first.")
        raise ValueError(
            f"Reaction '{reaction_id}': species '{species_id}' copy "
            f"{copy_index} not found. Available copies: {sorted(copies)}. "
            f"Use '{species_id}:{copy_index}' only if that copy was added "
            f"via addNode.")

    def _assign_glyph_ids(self):
        """Set the glyph_id field on every _NodeRecord."""
        counts: dict[str, int] = {}
        for node in self._nodes:
            counts[node.species_id] = counts.get(node.species_id, 0) + 1
        for node in self._nodes:
            node.glyph_id = f'sg_{node.species_id}_{node.copy_index}'

    def _glyph_id_for(self, species_id: str, copy_index: int) -> str:
        """Return the glyph id for a given species id + copy index."""
        for node in self._nodes:
            if node.species_id == species_id and node.copy_index == copy_index:
                return node.glyph_id
        raise KeyError(f"No node for {species_id}:{copy_index}")

    def _set_bb(self, obj, x: float, y: float, w: float, h: float):
        """Set bounding box on any SBML GraphicalObject."""
        bb = obj.getBoundingBox()
        bb.setX(x); bb.setY(y)
        bb.setWidth(w); bb.setHeight(h)

    def _build_compartment_glyphs(self, layout):
        for cid, rec in self._compartments.items():
            cg = layout.createCompartmentGlyph()
            cg.setId(f'cg_{cid}')
            cg.setCompartmentId(cid)
            self._set_bb(cg, rec.x, rec.y, rec.w, rec.h)

    def _build_species_glyphs(self, layout):
        """
        Create one SpeciesGlyph + one TextGlyph per node record.

        TextGlyphs are required by SBMLDiagrams — it crashes with an
        UnboundLocalError if the layout contains no TextGlyphs at all.
        The TextGlyph bounding box is set to match the node so that text
        renders centred without needing downstream fixup calls.
        """
        for node in self._nodes:
            sg = layout.createSpeciesGlyph()
            sg.setId(node.glyph_id)
            sg.setSpeciesId(node.species_id)
            self._set_bb(sg, node.x, node.y, node.w, node.h)

            tg = layout.createTextGlyph()
            tg.setId(f'tg_{node.glyph_id}')
            tg.setOriginOfTextId(node.species_id)
            tg.setGraphicalObjectId(node.glyph_id)
            self._set_bb(tg, node.x, node.y, node.w, node.h)

    def _build_reaction_glyphs(self, model, layout):
        """
        Create ReactionGlyphs with SpeciesReferenceGlyphs.

        Each SpeciesReferenceGlyph references the correct SpeciesGlyph copy
        via speciesGlyphId, ensuring alias nodes connect to the right arm.
        If bezier handles were supplied, CubicBezier curve segments are written
        into each arm's Curve element.
        """
        for rec in self._reactions:
            rid = rec.reaction_id
            rg = layout.createReactionGlyph()
            rg.setId(f'rg_{rid}_0')
            rg.setReactionId(rid)
            if rec.direct:
                # Direct reaction: use width=-1 as an unambiguous sentinel so
                # the renderer can reliably distinguish it from a junction at (0,0).
                cx, cy = rec.center if rec.center else (0, 0)
                self._set_bb(rg, cx, cy, -1, 0)
            elif rec.center:
                cx, cy = rec.center
                self._set_bb(rg, cx, cy, 0, 0)
            else:
                # No junction defined - compute centroid as default position.
                all_nodes = [self._get_node(sid, cidx)
                             for sid, cidx in rec.reactant_refs + rec.product_refs]
                if all_nodes:
                    cx = sum(n.x + n.w/2 for n in all_nodes) / len(all_nodes)
                    cy = sum(n.y + n.h/2 for n in all_nodes) / len(all_nodes)
                else:
                    cx, cy = 0, 0
                self._set_bb(rg, cx, cy, 0, 0)

            rxn = model.getReaction(rid)
            n_reactants = len(rec.reactant_refs)

            # -------------------------------------------------------------------
            # Direct reaction: single bezier from reactant to product, no junction
            # For a direct uni-uni with handles, store the full curve on the
            # reactant SpeciesReferenceGlyph (reactant->product via BP1/BP2).
            # -------------------------------------------------------------------
            if rec.direct:
                # Reactant glyph (no junction role — arm goes straight to product)
                for j, (sid, cidx) in enumerate(rec.reactant_refs):
                    srg = rg.createSpeciesReferenceGlyph()
                    srg.setId(f'srg_{rid}_r{j}')
                    srg.setSpeciesGlyphId(self._glyph_id_for(sid, cidx))
                    sr = rxn.getReactant(j)
                    if sr and sr.getId():
                        srg.setSpeciesReferenceId(sr.getId())
                    srg.setRole(libsbml.SPECIES_ROLE_SUBSTRATE)

                    # For a direct bezier, handles[j] gives (BP1, BP2) for the
                    # single curve from this reactant through to the product.
                    handle = rec.handles[j] if (rec.handles and j < len(rec.handles)) else None
                    if handle is not None and rec.product_refs:
                        rnode = self._get_node(sid, cidx)
                        pnode = self._get_node(*rec.product_refs[0])
                        rx = rnode.x + rnode.w / 2; ry = rnode.y + rnode.h / 2
                        px = pnode.x + pnode.w / 2; py = pnode.y + pnode.h / 2
                        (bp1x, bp1y), (bp2x, bp2y) = handle
                        self._add_cubic(srg, rx, ry, bp1x, bp1y, bp2x, bp2y, px, py)

                # Product glyph (role only, no curve — renderer uses reactant curve)
                for j, (sid, cidx) in enumerate(rec.product_refs):
                    srg = rg.createSpeciesReferenceGlyph()
                    srg.setId(f'srg_{rid}_p{j}')
                    srg.setSpeciesGlyphId(self._glyph_id_for(sid, cidx))
                    sr = rxn.getProduct(j)
                    if sr and sr.getId():
                        srg.setSpeciesReferenceId(sr.getId())
                    srg.setRole(libsbml.SPECIES_ROLE_PRODUCT)
                continue  # skip junction-based arm building below

            for j, (sid, cidx) in enumerate(rec.reactant_refs):
                srg = rg.createSpeciesReferenceGlyph()
                srg.setId(f'srg_{rid}_r{j}')
                srg.setSpeciesGlyphId(self._glyph_id_for(sid, cidx))
                sr = rxn.getReactant(j)
                if sr and sr.getId():
                    srg.setSpeciesReferenceId(sr.getId())
                srg.setRole(libsbml.SPECIES_ROLE_SUBSTRATE)

                # Write bezier curve if handles supplied for this arm
                handle = rec.handles[j] if (rec.handles and j < len(rec.handles)) else None
                if handle is not None:
                    node = self._get_node(sid, cidx)
                    (bp1x, bp1y), (bp2x, bp2y) = handle
                    ncx = node.x + node.w / 2
                    ncy = node.y + node.h / 2
                    # Store the bezier with the node CENTRE as start so the
                    # user's BP1/BP2 values (which are designed for centre-start)
                    # work correctly. The renderer clips the visible portion at
                    # the node's expanded bounding box edge.
                    self._add_cubic(srg, ncx, ncy, bp1x, bp1y, bp2x, bp2y, cx, cy)

            for j, (sid, cidx) in enumerate(rec.product_refs):
                srg = rg.createSpeciesReferenceGlyph()
                srg.setId(f'srg_{rid}_p{j}')
                srg.setSpeciesGlyphId(self._glyph_id_for(sid, cidx))
                sr = rxn.getProduct(j)
                if sr and sr.getId():
                    srg.setSpeciesReferenceId(sr.getId())
                srg.setRole(libsbml.SPECIES_ROLE_PRODUCT)

                # Write bezier curve if handles supplied for this arm
                arm_idx = n_reactants + j
                handle = rec.handles[arm_idx] if (rec.handles and arm_idx < len(rec.handles)) else None
                if handle is not None:
                    node = self._get_node(sid, cidx)
                    ncx = node.x + node.w / 2
                    ncy = node.y + node.h / 2
                    (bp1x, bp1y), (bp2x, bp2y) = handle
                    self._add_cubic(srg, cx, cy, bp1x, bp1y, bp2x, bp2y, ncx, ncy)

    def _add_cubic(self, srg, sx, sy, bp1x, bp1y, bp2x, bp2y, ex, ey):
        """Write a CubicBezier segment into a SpeciesReferenceGlyph's curve."""
        cb = srg.getCurve().createCubicBezier()
        cb.getStart().setX(sx);        cb.getStart().setY(sy)
        cb.getBasePoint1().setX(bp1x); cb.getBasePoint1().setY(bp1y)
        cb.getBasePoint2().setX(bp2x); cb.getBasePoint2().setY(bp2y)
        cb.getEnd().setX(ex);          cb.getEnd().setY(ey)

    @staticmethod
    def _node_edge_point_toward(cx: float, cy: float,
                                toward_x: float, toward_y: float,
                                rx: float, ry: float,
                                rw: float, rh: float) -> tuple:
        """
        Return the point where the ray from (cx,cy) toward (toward_x,toward_y)
        exits the rectangle (rx,ry,rw,rh).  Used to find where a bezier arm
        should start on the surface of an expanded node bounding box.
        """
        dx = toward_x - cx; dy = toward_y - cy
        if abs(dx) < 1e-9 and abs(dy) < 1e-9:
            return cx, cy
        candidates = []
        if abs(dx) > 1e-9:
            for ex in (rx, rx + rw):
                t = (ex - cx) / dx
                if t > 1e-9:
                    y = cy + t * dy
                    if ry <= y <= ry + rh:
                        candidates.append(t)
        if abs(dy) > 1e-9:
            for ey in (ry, ry + rh):
                t = (ey - cy) / dy
                if t > 1e-9:
                    x = cx + t * dx
                    if rx <= x <= rx + rw:
                        candidates.append(t)
        if not candidates:
            return cx, cy
        t_min = min(candidates)
        return cx + t_min * dx, cy + t_min * dy

    @staticmethod
    def _node_edge_point(node: '_NodeRecord', toward_x: float, toward_y: float,
                         gap: float = 4.0) -> tuple:
        """
        Return the point on the edge of node's bounding box on the line from
        the node centre toward (toward_x, toward_y), offset outward by gap pixels.
        """
        import math
        cx = node.x + node.w / 2
        cy = node.y + node.h / 2
        dx = toward_x - cx
        dy = toward_y - cy
        length = math.hypot(dx, dy)
        if length < 1e-6:
            return cx, cy

        # Find t for intersection with each edge
        candidates = []
        if abs(dx) > 1e-9:
            for ex in (node.x, node.x + node.w):
                t = (ex - cx) / dx
                if t > 1e-9:
                    y = cy + t * dy
                    if node.y <= y <= node.y + node.h:
                        candidates.append(t)
        if abs(dy) > 1e-9:
            for ey in (node.y, node.y + node.h):
                t = (ey - cy) / dy
                if t > 1e-9:
                    x = cx + t * dx
                    if node.x <= x <= node.x + node.w:
                        candidates.append(t)

        if not candidates:
            return cx, cy

        t_edge = min(candidates)
        # Add gap: move slightly beyond the edge along the direction
        t_gap = t_edge + gap / length
        return cx + t_gap * dx, cy + t_gap * dy

    def _get_node(self, species_id: str, copy_index: int) -> '_NodeRecord':
        """Return the _NodeRecord for species_id + copy_index."""
        for node in self._nodes:
            if node.species_id == species_id and node.copy_index == copy_index:
                return node
        raise KeyError(f"No node for {species_id}:{copy_index}")