import json
import os
import math
from functools import reduce

# PERSISTENCE FILE
REGISTRY_PATH = "./logos/memory/prime_registry.json"
os.makedirs(os.path.dirname(REGISTRY_PATH), exist_ok=True)

class PrimeTokenDB:
    def __init__(self):
        self.registry = self.load_registry()
        self.alias_registry = self.load_alias_registry() # New: URL/Sentence -> CompositeID Map
        try:
            # Handle potential string keys from JSON loading
            self.reverse_registry = {int(v) if isinstance(v, (int, str)) and str(v).isdigit() else v: k for k, v in self.registry.items()}
        except Exception as e:
             print(f"[PrimeDB] Warning rebuilding reverse registry: {e}")
             self.reverse_registry = {}
        
        # Protocol 33: Load Genesis Block BEFORE calculating next prime
        self.load_genesis_block()
             
        self.next_prime = self.calculate_next_prime()

    def load_registry(self):
        if os.path.exists(REGISTRY_PATH):
            try:
                with open(REGISTRY_PATH, 'r') as f:
                    data = json.load(f)
                    # Convert values to int just in case
                    return {k: int(v) for k, v in data.items()}
            except:
                return {"_ROOT": 1}
        return {"_ROOT": 1} # Identity

    def load_alias_registry(self):
        """Loads the Semantic Alias Map (URLs -> Prime Composites)"""
        alias_path = REGISTRY_PATH.replace('prime_registry.json', 'alias_registry.json')
        if os.path.exists(alias_path):
             try:
                 with open(alias_path, 'r') as f:
                     return json.load(f)
             except: return {}
        return {}

    def save_registry(self):
        with open(REGISTRY_PATH, 'w') as f:
            json.dump(self.registry, f, indent=2)
            
        alias_path = REGISTRY_PATH.replace('prime_registry.json', 'alias_registry.json')
        with open(alias_path, 'w') as f:
            json.dump(self.alias_registry, f, indent=2)

    def register_alias(self, source_string, composite_id):
        """
        Maps a complex string (URL, Sentence) to an existing Composite ID.
        This allows 'Source Attribution' - knowing a string implies a specific logic state.
        """
        self.alias_registry[source_string] = composite_id
        self.save_registry()

    def resolve_alias(self, text):
        """
        Checks if text matches a known alias. Returns Composite ID or None.
        """
        return self.alias_registry.get(text)

    def is_prime(self, n):
        """Standard Primality Test"""
        if n < 2: return False
        for i in range(2, int(math.sqrt(n)) + 1):
            if n % i == 0: return False
        return True

    def load_genesis_block(self):
        """
        Protocol 33: Reserve Genesis Axioms
        Ensures Primes [2..29] are reserved for System Concepts.
        """
        try:
            mtl_dir = os.path.dirname(os.path.dirname(__file__)) # .../logos
            genesis_path = os.path.join(mtl_dir, "mtl", "genesis.json")
            if os.path.exists(genesis_path):
                with open(genesis_path, 'r') as f:
                    genesis = json.load(f)
                    for p_str, token in genesis.items():
                        p = int(p_str)
                        self.registry[token] = p
                        self.reverse_registry[p] = token
        except Exception as e:
            print(f"[PrimeDB] Genesis Load Error: {e}")

    def calculate_next_prime(self):
        """Finds the next available prime slot for a new token."""
        # Check Genesis first
        self.load_genesis_block()
        
        if not self.registry: return 31 # Start after Genesis
        
        # Ensure we look at integer values
        values = [int(v) for v in self.registry.values()]
        highest = max(values) if values else 1
        
        # Protocol 33: Hard Floor at 30
        if highest < 30: highest = 30
            
        candidate = highest + 1
        while not self.is_prime(candidate):
            candidate += 1
        return candidate

    def get_token_prime(self, token):
        """
        Retrieves the Prime Coordinate for a token.
        If new, assigns the next prime and expands the Manifold.
        """
        token = str(token).lower().strip()
        if token in self.registry:
            return self.registry[token]
        
        # MINT NEW PRIME
        prime = self.next_prime
        self.registry[token] = prime
        self.reverse_registry[prime] = token
        
        # Calc next for future
        self.next_prime += 1
        while not self.is_prime(self.next_prime):
            self.next_prime += 1
            
        self.save_registry()
        return prime

    def encode_state(self, tokens, source_ref=None, meta_tensor=None):
        """
        Converts a list of concepts (text) into a Single Unique Integer.
        Example: ["AI", "Logic"] -> 3 * 5 = 15
        
        source_ref: Optional string (URL) to alias this state to.
        meta_tensor: Optional dictionary containing Vector/Gradient data.
        """
        primes = [self.get_token_prime(t) for t in tokens]
        # Calculate Product (The Composite Weight)
        if not primes: return 1, []
        composite = reduce(lambda x, y: x * y, primes, 1)
        
        if source_ref:
            self.register_alias(source_ref, composite)
        
        # If meta_tensor provided, we might want to store it or return it tied to the composite
        # For now, the system expects (composite, primes), so we maintain that contract
        # But we stash the meta data in the alias registry for now as a hack or return it in a wrapped object?
        # Let's keep the return signature simple for now to not break the Atomizer.
            
        return composite, primes

    def decode_state(self, composite_integer):
        """
        Reconstructs the atoms from the composite weight.
        (Factoring the Manifold State)
        """
        factors = []
        # Optimization: We only check primes we know exist in the registry
        known_primes = sorted([int(k) for k in self.reverse_registry.keys()])
        
        temp = composite_integer
        for p in known_primes:
            if p <= 1: continue
            if p * p > temp: 
                # If what remains is prime and in registry
                if temp in self.reverse_registry:
                     factors.append(self.reverse_registry[temp])
                     temp = 1
                break 
            
            while temp % p == 0:
                factors.append(self.reverse_registry[p])
                temp //= p
        
        if temp > 1 and temp in self.reverse_registry:
            factors.append(self.reverse_registry[temp])
            
        return factors