# Copyright 2026 Gaofeng Fan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools
import inspect
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Union, Callable, ClassVar


# Global name counter shared across Circuit and Testbench
# This prevents name collisions when both have the same base name
_global_name_counters: Dict[str, int] = {}


def _get_unique_name(base_name: str) -> str:
    """
    Get a unique name using the global counter.

    First use of a name returns it as-is.
    Subsequent uses return name_1, name_2, etc.
    """
    if base_name in _global_name_counters:
        _global_name_counters[base_name] += 1
        return f"{base_name}_{_global_name_counters[base_name]}"
    else:
        _global_name_counters[base_name] = 0
        return base_name


def reset_global_name_counters():
    """Reset the global name counters. Useful for testing."""
    _global_name_counters.clear()


def _get_assignment_target() -> Optional[str]:
    """
    Try to capture the variable name from the caller's source code.

    This inspects the call stack to find the line where a Circuit is being
    assigned to a variable, and extracts the variable name.

    Returns:
        The variable name if found, None otherwise.

    Note:
        This works for most cases but fails in:
        - REPL/Jupyter (no source file)
        - list.append(Circuit(...)) style calls
        - Complex expressions
    """
    try:
        for frame_info in inspect.stack()[3:]:  # Skip this func, __init__, and immediate caller
            if frame_info.code_context:
                line = frame_info.code_context[0]
                # Skip decorator wrapper internals (e.g., "result = func(*args, **kwargs)")
                if 'func(*args' in line or 'func( *args' in line:
                    continue
                # Match: var_name = func_name( or var_name = Class(
                match = re.match(r'\s*(\w+)\s*=\s*\w+\s*\(', line)
                if match:
                    var_name = match.group(1)
                    # Skip common temporary variable names used in wrappers/frameworks
                    if var_name not in ('result', 'ret', 'rv', 'value', '_', 'code', 'obj', 'item', 'x'):
                        return var_name
    except Exception:
        pass
    return None


@dataclass
class Net:
    name: str


@dataclass(frozen=True)
class ParamRef:
    """Reference to a subcircuit parameter for netlist emission."""
    name: str

    def __str__(self) -> str:
        return self.name


def param(name: str) -> ParamRef:
    """Helper to create a ParamRef by name."""
    return ParamRef(name)

@dataclass
class Instance:
    name: str
    kind: str                 # "mos", "res", "cap", ...
    pins: Dict[str, Net]     # {"d": Net, "g": Net, ...}
    params: Dict[str, float] # {"w": 1e-6, "l": 180e-9}
    model: str = None        # model name (optional)





@dataclass
class SubcircuitInstance:
    """Instance of a subcircuit within a parent circuit

    Schematic positioning supports both absolute and relative modes:
        - Absolute: schematic_position=(x, y)
        - Relative: schematic_position={'relative_to': 'I_OLED', 'x_shift': -2, 'y_shift': 0}

    Relative positioning is resolved at schematic generation time.
    """
    name: str                           # instance name, e.g., "X1"
    subcircuit: 'Subcircuit'           # reference to subcircuit definition
    connections: Dict[str, Net]         # port_name -> connected net
    params: Dict[str, any] = field(default_factory=dict)  # parameter values
    param_blacklist: List[str] = field(default_factory=list)  # params to hide in drawings
    schematic_position: Optional[Union[Tuple[float, float], Dict[str, any]]] = None  # (x, y) or relative spec


@dataclass
class _CircuitBase:
    """Internal base class for circuit containers."""
    name: str
    nets: Dict[str, Net] = field(default_factory=dict)
    instances: List[Instance] = field(default_factory=list)

    subcircuit_instances: List['CircuitInstance'] = field(default_factory=list)
    _instance_counters: Dict[str, int] = field(default_factory=dict, repr=False)

    def _generate_instance_name(self, circuit_name: str) -> str:
        """
        Generate a unique instance name based on circuit name.

        Pattern: I_{circuit_name} for first instance, I_{circuit_name}_{n} for subsequent.
        Example: I_oled_blue, I_oled_blue_1, I_oled_blue_2
        """
        base_name = f"I_{circuit_name}"
        if base_name not in self._instance_counters:
            self._instance_counters[base_name] = 0
            return base_name
        else:
            self._instance_counters[base_name] += 1
            return f"{base_name}_{self._instance_counters[base_name]}"

    def net(self, name_or_net: Union[str, Net]) -> Net:
        """
        Get or create a net by name, or return a Net object directly.

        Args:
            name_or_net: Net name string or Net object

        Returns:
            Net object
        """
        if isinstance(name_or_net, Net):
            return name_or_net

        if name_or_net not in self.nets:
            self.nets[name_or_net] = Net(name_or_net)
        return self.nets[name_or_net]

    def add(self, inst: Instance):
        self.instances.append(inst)
        return inst

    def gnd(self):
        return self.net("0")

    def vdd(self, name="vdd"):
        return self.net(name)

    def vss(self, name="vss"):
        return self.net(name)



    def instantiate(self, circuit: 'Circuit', instance_name: str = None,
                    schematic_position: Tuple[float, float] = None,
                    **kwargs) -> 'CircuitInstance':
        """
        Instantiate a circuit within this circuit.

        Automatically separates port connections from parameters based on
        the circuit's port definitions.

        Args:
            circuit: The Circuit to instantiate
            instance_name: Name for this instance (e.g., "X1", "R1", "I_OLED").
                          If None, auto-generates: I_{circuit_name}, I_{circuit_name}_1, etc.
            schematic_position: Optional (x, y) tuple for manual placement in schematic view.
                               If None, automatic grid placement is used.
            **kwargs: Port connections and parameters mixed together.
                     - If kwarg name is in circuit.ports → connection
                     - Otherwise → parameter

        Returns:
            CircuitInstance

        Example:
            # Built-in device with parameters:
            my_circuit.instantiate(resistor, "R1", p=vdd, n=gnd, r=1000)
            # p, n are ports (connections), r is parameter

            # User-defined circuit:
            my_circuit.instantiate(oled_cell, "XOLED1", ANODE=v_single, ELVSS=gnd)
            # ANODE, ELVSS are ports (connections)

            # With Verilog-A parameters:
            oled_lut = generic("oled_LUT_dc", ["p", "n"])
            my_circuit.instantiate(oled_lut, "I1", p=anode, n=cathode,
                                   tmdata="model.txt", area=938.488)
            # p, n are ports; tmdata, area are parameters

            # With manual schematic position:
            my_circuit.instantiate(cap_lut, "C1", p=anode, n=cathode,
                                   schematic_position=(5.0, -3.0))
        """
        # Auto-generate instance name if not provided
        if instance_name is None:
            instance_name = self._generate_instance_name(circuit.name)

        # Separate port connections from parameters
        port_connections = {}
        params = {}

        for key, value in kwargs.items():
            if key in circuit.ports:
                port_connections[key] = value
            else:
                params[key] = value

        # Validate all ports are connected
        for port in circuit.ports:
            if port not in port_connections:
                raise ValueError(
                    f"Port '{port}' of circuit '{circuit.name}' not connected "
                    f"in instance '{instance_name}'. "
                    f"Required ports: {circuit.ports}"
                )

        # Convert string connections to Net objects
        connections = {}
        for port, net in port_connections.items():
            if isinstance(net, str):
                connections[port] = self.net(net)
            elif isinstance(net, Net):
                connections[port] = net
            else:
                raise TypeError(f"Connection for port '{port}' must be Net or str, got {type(net)}")

        inst = CircuitInstance(
            name=instance_name,
            subcircuit=circuit,
            connections=connections,
            params=params,
            schematic_position=schematic_position
        )
        self.subcircuit_instances.append(inst)
        return inst

    def add_instance(self, circuit: 'Circuit', instance_name: str = None,
                     schematic_position: Tuple[float, float] = None,
                     **kwargs) -> 'CircuitInstance':
        """Alias for instantiate() with a friendlier name."""
        return self.instantiate(
            circuit,
            instance_name=instance_name,
            schematic_position=schematic_position,
            **kwargs,
        )


@dataclass
class Circuit(_CircuitBase):
    """
    A reusable circuit block with defined ports.

    This is the primary building block in analogpy. It maps to SPICE/Spectre's
    'subckt' keyword. Use Circuit to define reusable blocks that can be
    instantiated multiple times.

    Circuit names are automatically made unique:
    1. First, tries to capture the Python variable name (e.g., oled_single_blue)
    2. Falls back to auto-incrementing (e.g., oled_model, oled_model_1, oled_model_2)

    Aliases: Subcircuit, Subckt (all equivalent)

    Port Directions:
        Ports can optionally have directions for schematic visualization:
        - "input": Triangle pointing INTO the symbol
        - "output": Triangle pointing OUT of the symbol
        - "inout": Diamond shape (default if not specified)

        Specify directions using colon syntax in port names:
            Circuit("amp", ports=["VIN:input", "VOUT:output", "VDD", "VSS"])
        Ports without colon default to "inout".

    Schematic Wires:
        By default, nets use labels (implicit connections). To draw explicit
        wires for specific nets, use the wired_nets() method:
            oled.wired_nets('LUM', 'ANODE')  # Draw wires for these nets

    Example:
        # Simple ports (all default to inout):
        oled = Circuit("oled_cell", ports=["ANODE", "ELVSS"])

        # With explicit directions using colon syntax:
        amp = Circuit("amplifier", ports=["IN:input", "OUT:output", "VDD", "VSS"])

        # Names are auto-captured from variable names:
        oled_blue = Circuit("oled", ports=["A", "B"])   # name = "oled_blue"
        oled_red = Circuit("oled", ports=["A", "B"])    # name = "oled_red"
    """
    ports: List[str] = field(default_factory=list)
    port_directions: Dict[str, str] = field(default_factory=dict)
    parameters: List[str] = field(default_factory=list)
    comments: List[str] = field(default_factory=list)
    _wired_nets: List[str] = field(default_factory=list, repr=False)
    _wired_net_ports: Dict[str, List[Tuple[str, str]]] = field(default_factory=dict, repr=False)

    # Valid port direction values
    VALID_PORT_DIRECTIONS: ClassVar[tuple] = ("input", "output", "inout")

    def __init__(self, name: str, ports: List[str], parameters: List[str] = None):
        # Hybrid naming strategy:
        # - First use of a base name: use as-is (respect explicit name)
        # - Subsequent uses: try to capture variable name, fall back to global counter
        if name in _global_name_counters:
            # Name already used - must resolve conflict
            var_name = _get_assignment_target()
            if var_name and var_name != name and var_name not in _global_name_counters:
                # Use variable name if it's unique
                final_name = var_name
                _global_name_counters[var_name] = 0
            else:
                # Fall back to global counter
                final_name = _get_unique_name(name)
        else:
            # First use of this base name - register it
            final_name = _get_unique_name(name)

        super().__init__(final_name)

        # Parse ports with optional direction syntax (e.g., "IN:input", "OUT:output")
        parsed_ports = []
        self.port_directions = {}
        for port_spec in ports:
            if ":" in port_spec:
                port_name, direction = port_spec.split(":", 1)
                if direction not in self.VALID_PORT_DIRECTIONS:
                    raise ValueError(
                        f"Invalid direction '{direction}' for port '{port_name}'. "
                        f"Valid values: {self.VALID_PORT_DIRECTIONS}"
                    )
                parsed_ports.append(port_name)
                self.port_directions[port_name] = direction
            else:
                parsed_ports.append(port_spec)
                self.port_directions[port_spec] = "inout"  # Default

        self.ports = parsed_ports
        self.parameters = parameters or []
        self.comments = []
        self._wired_nets = []  # Nets that should have explicit wires in schematic
        self._wired_net_ports = {}  # Net -> list of (instance_name, port_name) pairs
        self.param_blacklist = []  # Params to hide in drawings for this cell type

        # Capture the module where this instance was created for auto-discovery
        # of SYMBOL_CONFIGS in visualization.
        try:
            # Skip frames to find the actual user module
            for frame_info in inspect.stack():
                if 'analogpy/circuit.py' in frame_info.filename or 'analogpy\\circuit.py' in frame_info.filename:
                    continue
                self._module_name = inspect.getmodule(frame_info.frame).__name__
                break
        except Exception:
            self._module_name = None

        # Pre-create nets for all ports
        for port in self.ports:
            self.net(port)

    def get_port_direction(self, port_name: str) -> str:
        """
        Get the direction of a port.

        Args:
            port_name: Name of the port

        Returns:
            Direction string: "input", "output", or "inout" (default)
        """
        return self.port_directions.get(port_name, "inout")

    def comment(self, text: str) -> 'Circuit':
        """Add a comment to this circuit for documentation/traceability."""
        self.comments.append(text)
        return self

    def draw_wires(self, *net_names: Union[str, Net], only: Optional[List[Tuple[str, str]]] = None) -> 'Circuit':
        """
        Specify nets that should have explicit wires drawn in schematic view.

        By default, nets use labels (implicit connections). Use this method
        to draw explicit wires for specific nets.

        Args:
            *net_names: Net names or Net objects to draw wires for (e.g., 'LUM', mid)
            only: Optional list of (instance_name, port_name) pairs. If provided,
                  wires are drawn only between those specific ports for the given net.

        Returns:
            self for chaining
        """
        # Normalize all to strings
        names = [n.name if isinstance(n, Net) else n for n in net_names]

        if only is not None:
            if len(names) != 1:
                raise ValueError("draw_wires(..., only=...) supports exactly one net name.")
            net = names[0]
            if net not in self._wired_net_ports:
                self._wired_net_ports[net] = []
            self._wired_net_ports[net].extend(list(only))
            if net not in self._wired_nets:
                self._wired_nets.append(net)
            return self

        self._wired_nets.extend(names)
        return self

    def get_wired_nets(self) -> List[str]:
        """Get list of nets that should have explicit wires drawn."""
        return self._wired_nets

    def get_wired_net_ports(self) -> Dict[str, List[Tuple[str, str]]]:
        """Get mapping of net -> list of (instance_name, port_name) for selective wiring."""
        return self._wired_net_ports

    @classmethod
    def reset_name_counters(cls):
        """Reset the name counters. Useful for testing."""
        reset_global_name_counters()


# Aliases for compatibility and convenience
Subcircuit = Circuit  # Backwards compatibility
Subckt = Circuit      # SPICE terminology


# Also rename SubcircuitInstance to CircuitInstance for consistency
CircuitInstance = SubcircuitInstance


def auto_comment_params(func: Callable[..., Circuit]) -> Callable[..., Circuit]:
    """
    Decorator that automatically adds function arguments as comments to the returned Circuit.

    Use this on factory functions that create Circuit objects. All function arguments
    will be captured and added as comments to the Circuit for traceability.

    Example:
        @auto_comment_params
        def make_oled_circuit(oled_model_d1: str, oled_size: float) -> Circuit:
            oled = Circuit("oled_model", ports=["ANODE", "ELVSS"])
            # ... build circuit ...
            return oled

        # When called:
        circuit = make_oled_circuit(oled_model_d1="model.txt", oled_size=938.488)
        # The circuit will automatically have comments:
        #   oled_model_d1 = model.txt
        #   oled_size = 938.488
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # Get function signature to map positional args to parameter names
        sig = inspect.signature(func)
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()

        # Call the original function
        result = func(*args, **kwargs)

        # Add comments for all arguments if result is a Circuit
        if isinstance(result, Circuit):
            for param_name, value in bound.arguments.items():
                result.comment(f"{param_name} = {value}")

        return result

    return wrapper
