# 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.

"""Tests for SaveConfig."""

from analogpy import (
    SaveConfig, Testbench, Circuit, nmos, pmos,
    generate_spectre, DC
)
from analogpy.devices import vsource


def test_save_config_basic():
    """Test basic SaveConfig creation."""
    saves = SaveConfig("test")
    saves.voltage("out", "bias")
    saves.current("M1:d")
    saves.op("M1:gm", "M1:vth")

    assert len(saves) == 5
    assert "out" in saves.signals
    assert "M1:d" in saves.currents
    assert "M1:gm" in saves.op_params


def test_save_config_chaining():
    """Test fluent interface."""
    saves = (SaveConfig("ota")
             .voltage("out", "tail")
             .current("M1:d", "M2:d")
             .op("M1:gm"))

    assert len(saves) == 5


def test_save_config_prefix():
    """Test hierarchical prefix."""
    saves = SaveConfig("ota").voltage("out", "tail").op("M1:gm")
    prefixed = saves.with_prefix("X_LDO.X_OTA")

    assert "X_LDO.X_OTA.out" in prefixed.signals
    assert "X_LDO.X_OTA.tail" in prefixed.signals
    assert "X_LDO.X_OTA.M1:gm" in prefixed.op_params


def test_save_config_merge():
    """Test merging SaveConfigs."""
    saves1 = SaveConfig("a").voltage("out1")
    saves2 = SaveConfig("b").voltage("out2").current("M1:d")

    saves1.merge(saves2)

    assert "out1" in saves1.signals
    assert "out2" in saves1.signals
    assert "M1:d" in saves1.currents


def test_save_config_tags():
    """Test tagged signals."""
    saves = (SaveConfig("ota")
             .voltage("out", tag="essential")
             .voltage("tail", "bias", tag="internal")
             .op("M1:gm", "M2:gm", tag="op_params"))

    # Only essential
    essential = saves.only_tags("essential")
    assert "out" in essential.signals
    assert "tail" not in essential.signals

    # Without op_params
    no_op = saves.without_tags("op_params")
    assert "out" in no_op.signals
    assert len(no_op.op_params) == 0


def test_save_config_spectre_output():
    """Test Spectre output generation."""
    saves = (SaveConfig()
             .voltage("X_AMP.out", "X_AMP.bias")
             .current("X_AMP.M1:d")
             .op("X_AMP.M1:gm"))

    spectre = saves.to_spectre()
    assert "save X_AMP.out X_AMP.bias" in spectre
    assert "save X_AMP.M1:d" in spectre
    assert "save X_AMP.M1:gm" in spectre


def test_save_config_ngspice_output():
    """Test ngspice output generation."""
    saves = SaveConfig().voltage("out", "bias")
    ngspice = saves.to_ngspice()

    assert ".save v(out)" in ngspice
    assert ".save v(bias)" in ngspice


def test_testbench_save_integration():
    """Test SaveConfig integration with Testbench."""
    # Define OTA saves
    ota_saves = (SaveConfig("ota")
                 .voltage("out", "tail")
                 .op("M1:gm"))

    # Create simple inverter as DUT
    inv = Circuit("inverter", ports=["in", "out", "vdd", "vss"])
    inv.add_instance(nmos, "MN", d=inv.net("out"), g=inv.net("in"),
                    s=inv.net("vss"), b=inv.net("vss"), w=1e-6, l=180e-9)
    inv.add_instance(pmos, "MP", d=inv.net("out"), g=inv.net("in"),
                    s=inv.net("vdd"), b=inv.net("vdd"), w=2e-6, l=180e-9)

    # Create testbench
    tb = Testbench("tb_test")
    vdd = tb.net("vdd")
    gnd = tb.gnd()  # Global ground "0" at testbench level
    tb.add_instance(vsource, instance_name="I_vdd", p=vdd, n=gnd, dc=1.8)
    tb.add_instance(inv, "X_INV",
                   **{"in": tb.net("in"), "out": tb.net("out"),
                      "vdd": vdd, "vss": gnd})

    # Add saves with prefix
    tb.save(ota_saves.with_prefix("X_INV"))
    tb.save("vin")  # Direct signal

    tb.add_analysis(DC())

    netlist = generate_spectre(tb)
    print("\n=== Testbench with Saves ===")
    print(netlist)

    assert "save X_INV.out X_INV.tail" in netlist
    assert "vin" in netlist  # vin is merged into voltage saves


def test_save_only_override():
    """Test save_only overrides all previous saves."""
    tb = Testbench("tb_test")

    # Add some saves
    tb.save(SaveConfig().voltage("a", "b", "c"))
    tb.save("d")

    # Override with only specific signals
    tb.save_only("x", "y")

    all_saves = tb.get_all_saves()
    assert "x" in all_saves.signals
    assert "y" in all_saves.signals
    assert "a" not in all_saves.signals


def test_nested_hierarchy_saves():
    """Test saves with nested hierarchy."""
    # OTA saves
    ota_saves = SaveConfig("ota").voltage("out").op("M1:gm")

    # LDO saves (includes OTA)
    ldo_saves = SaveConfig("ldo").voltage("vout", "gate")
    ldo_full_saves = SaveConfig("ldo_full")
    ldo_full_saves.merge(ldo_saves)
    ldo_full_saves.merge(ota_saves.with_prefix("X_OTA"))

    # Apply to testbench level
    final_saves = ldo_full_saves.with_prefix("X_LDO")

    assert "X_LDO.vout" in final_saves.signals
    assert "X_LDO.X_OTA.out" in final_saves.signals
    assert "X_LDO.X_OTA.M1:gm" in final_saves.op_params

    spectre = final_saves.to_spectre()
    print("\n=== Nested Hierarchy Saves ===")
    print(spectre)


if __name__ == "__main__":
    test_save_config_basic()
    print("test_save_config_basic passed")

    test_save_config_chaining()
    print("test_save_config_chaining passed")

    test_save_config_prefix()
    print("test_save_config_prefix passed")

    test_save_config_merge()
    print("test_save_config_merge passed")

    test_save_config_tags()
    print("test_save_config_tags passed")

    test_save_config_spectre_output()
    print("test_save_config_spectre_output passed")

    test_save_config_ngspice_output()
    print("test_save_config_ngspice_output passed")

    test_testbench_save_integration()
    print("test_testbench_save_integration passed")

    test_save_only_override()
    print("test_save_only_override passed")

    test_nested_hierarchy_saves()
    print("test_nested_hierarchy_saves passed")

    print("\nAll save tests passed!")
