CATEGORY_ORDER = [
    "CPU", "Motherboard", "CPU Cooler", "Memory",
    "Storage", "Video Card", "Case", "Power Supply"
]

import random

class CSPBacktracking:
    def __init__(self, components):
        self.components = components
        self.by_cat = {cat: [c for c in components if c["category"] == cat]
                       for cat in CATEGORY_ORDER}
        
        # Shuffle components for variety
        for cat in CATEGORY_ORDER:
            random.shuffle(self.by_cat[cat])
        
        # Log component counts
        import logging
        for cat in CATEGORY_ORDER:
            count = len(self.by_cat.get(cat, []))
            if count > 0:
                logging.info("CSP: Category '%s' has %d components", cat, count)
            else:
                logging.warning("CSP: Category '%s' has NO components!", cat)

    def calculate_power_requirement(self, partial_build):
        total = 0
        for cat, comp in partial_build.items():
            if not comp or cat == "Power Supply":
                continue
            tdp = comp.get("attrs", {}).get("tdp", 0) or comp.get("attrs", {}).get("powerRequirement", 0)
            total += tdp
        
        if total == 0:
            defaults = {
                "CPU": 65,
                "Video Card": 150,
                "Memory": 5,
                "Storage": 5,
                "Motherboard": 30,
                "CPU Cooler": 5,
            }
            for cat, comp in partial_build.items():
                if comp and cat in defaults:
                    total += defaults[cat]
        
        return int(total * 1.2) if total > 0 else 400

    def is_compatible(self, partial_build, new_component):
        cat = new_component["category"]
        cpu = partial_build.get("CPU")
        mobo = partial_build.get("Motherboard")
        cooler = partial_build.get("CPU Cooler")
        ram = partial_build.get("Memory")
        gpu = partial_build.get("Video Card")
        pc_case = partial_build.get("Case")
        psu = partial_build.get("Power Supply")

        def get_attr(comp, key, default=None):
            if not comp:
                return default
            attrs = comp.get("attrs") or {}
            value = attrs.get(key, default)
            return value if value is not None else default

        # CPU ↔ Motherboard
        if cat == "Motherboard" and cpu:
            cpu_socket = get_attr(cpu, "socket")
            mobo_socket = get_attr(new_component, "socket")
            if cpu_socket and mobo_socket and cpu_socket != mobo_socket:
                return False
        if cat == "CPU" and mobo:
            cpu_socket = get_attr(new_component, "socket")
            mobo_socket = get_attr(mobo, "socket")
            if cpu_socket and mobo_socket and cpu_socket != mobo_socket:
                return False

        # CPU Cooler ↔ CPU socket
        if cat == "CPU Cooler" and cpu:
            cpu_socket = get_attr(cpu, "socket")
            supported = get_attr(new_component, "supported_sockets", [])
            if cpu_socket and supported and cpu_socket not in supported:
                return False
        if cat == "CPU" and cooler:
            cpu_socket = get_attr(new_component, "socket")
            supported = get_attr(cooler, "supported_sockets", [])
            if cpu_socket and supported and cpu_socket not in supported:
                return False

        # Memory ↔ Motherboard
        if cat == "Memory" and mobo:
            ram_type = get_attr(new_component, "ram_type")
            mobo_ram_type = get_attr(mobo, "ram_type")
            if ram_type and mobo_ram_type and ram_type != mobo_ram_type:
                return False
        if cat == "Motherboard" and ram:
            ram_type = get_attr(ram, "ram_type")
            mobo_ram_type = get_attr(new_component, "ram_type")
            if ram_type and mobo_ram_type and ram_type != mobo_ram_type:
                return False

        # PSU ↔ Total Power
        if cat == "Power Supply":
            psu_wattage = get_attr(new_component, "wattage", 0)
            power_req = self.calculate_power_requirement(partial_build)
            if psu_wattage > 0 and power_req > 0 and psu_wattage < power_req:
                return False

        # GPU ↔ Case
        if cat == "Video Card" and pc_case:
            gpu_length = get_attr(new_component, "length", 0)
            max_length = get_attr(pc_case, "maxGpuLength", 0)
            if max_length > 0 and gpu_length > 0 and gpu_length > max_length:
                return False
        if cat == "Case" and gpu:
            gpu_length = get_attr(gpu, "length", 0)
            max_length = get_attr(new_component, "maxGpuLength", 0)
            if max_length > 0 and gpu_length > 0 and gpu_length > max_length:
                return False

        # Cooler ↔ Case height
        if cat == "CPU Cooler" and pc_case:
            cooler_height = get_attr(new_component, "height", 0)
            max_height = get_attr(pc_case, "maxCoolerHeight", 0)
            if max_height > 0 and cooler_height > 0 and cooler_height > max_height:
                return False
        if cat == "Case" and cooler:
            cooler_height = get_attr(cooler, "height", 0)
            max_height = get_attr(new_component, "maxCoolerHeight", 0)
            if max_height > 0 and cooler_height > 0 and cooler_height > max_height:
                return False

        return True

    def backtrack(self, categories, partial_build, budget):
        current_cost = sum(c["price"] for c in partial_build.values() if c)
        if current_cost > budget:
            return
        
        if not categories:
            if current_cost <= budget:
                yield partial_build.copy()
            return

        next_cat = categories[0]

        if next_cat in partial_build:
            yield from self.backtrack(categories[1:], partial_build, budget)
            return

        available_components = self.by_cat.get(next_cat, [])
        available_components = [c for c in available_components if c.get("price", 0) > 0]
        
        if not available_components:
            return
        
        # Try components in order (already shuffled for variety)
        for comp in available_components:
            comp_price = comp.get("price", 0)
            if current_cost + comp_price > budget:
                continue
            
            if self.is_compatible(partial_build, comp):
                partial_build[next_cat] = comp
                yield from self.backtrack(categories[1:], partial_build, budget)
                del partial_build[next_cat]

    def solve(self, budget, user_inputs={}):
        partial_build = {}
        id_map = {c["id"]: c for c in self.components}

        for cat, comp_id in user_inputs.items():
            if comp_id in id_map:
                partial_build[cat] = id_map[comp_id]

        pre_selected_cost = sum(c["price"] for c in partial_build.values() if c)
        if pre_selected_cost > budget:
            return

        remaining_cats = [c for c in CATEGORY_ORDER if c not in partial_build]
        seen_solutions = set()
        
        for solution in self.backtrack(remaining_cats, partial_build, budget):
            solution_signature = tuple(sorted(
                (cat, comp["id"]) for cat, comp in solution.items() if comp
            ))
            
            if solution_signature not in seen_solutions:
                seen_solutions.add(solution_signature)
                yield solution