import unittest
from wif2ansible.models import Server, Flow
from wif2ansible.inventory import generate_inventory

class TestInventoryKeys(unittest.TestCase):
    def test_inventory_keys_are_hostnames(self):
        # Create a server with Ref, Hostname, IP
        s1 = Server(reference="SERVER_REF_01", hostname="server01", ip_addresses=["192.168.1.10"], platform="windows")
        
        # Create a flow matching this server
        f1 = Flow(flow_id="1", source_ip="192.168.1.10", destination_ip="10.0.0.1", ports=[80])
        
        servers = {"SERVER_REF_01": s1}
        flows = [f1]
        
        inventory = generate_inventory(servers, flows)
        
        # Verify stricture
        hosts = inventory['all']['hosts']
        
        # Key should be HOSTNAME "server01" (prioritized over Ref)
        self.assertIn("server01", hosts)
        self.assertNotIn("192.168.1.10", hosts)
        
        # Check variables
        host_vars = hosts["server01"]
        self.assertEqual(host_vars['ansible_host'], "192.168.1.10")
        self.assertEqual(host_vars['ansible_connection'], "winrm")

    def test_clean_reference_logic(self):
        from wif2ansible.parsers import clean_reference
        
        # Test cases
        self.assertEqual(clean_reference("SRV123 MyServer"), "MyServer")
        self.assertEqual(clean_reference("SVR999 AnotherServer"), "AnotherServer")
        self.assertEqual(clean_reference("srv001 lowercase"), "lowercase")
        self.assertEqual(clean_reference("SvR555 MixedCase"), "MixedCase")
        self.assertEqual(clean_reference("JustName"), "JustName")
        self.assertEqual(clean_reference("SRV123"), "") # Should be empty? or handle? 
        # If it's just SRV123, strip returns empty. 
        # User said "never include these in output". 
        # If the server is ONLY named SRV123, what then? 
        # Assuming there is usually a name. 
        
if __name__ == '__main__':
    unittest.main()