from collections import defaultdict
import random
from typing import Optional
from uuid import UUID

from sqlalchemy import Delete, Select

from src.auth.service import CurrentUser
from src.contribution_util import calculate_contribution_accurate
from src.database.core import CollectorDbSession, DbSession
# from src.scope_equipment.model import ScopeEquipment
# from src.scope_equipment.service import get_by_scope_name
from src.overhaul_activity.service import get_all_by_session_id, get_standard_scope_by_session_id

# async def get_all_budget_constrains(
#     *, db_session: DbSession, session_id: str, cost_threshold: float = 100000000
# ):

# At the module level, add this dictionary to store persistent EAF values
from collections import defaultdict
from uuid import UUID
from typing import List, Dict, Tuple

from src.database.core import CollectorDbSession, DbSession
from src.overhaul_activity.service import get_standard_scope_by_session_id
from src.contribution_util import calculate_contribution_accurate


async def get_all_budget_constrains(
    *,
    db_session: DbSession,
    collector_db: CollectorDbSession,
    session_id: UUID,
    simulation_result: Dict,
    cost_threshold: float = 100_000_000,
    use_optimal: bool = True,   # default = optimal (knapsack)
) -> Tuple[List[Dict], List[Dict]]:
    """
    Select equipment under budget constraint using contribution + cost efficiency.
    Returns (priority_list, consequence_list).
    """

    calc_result = simulation_result["calc_result"]
    plant_result = simulation_result["plant_result"]

    equipments = await get_standard_scope_by_session_id(
        db_session=db_session,
        overhaul_session_id=session_id,
        collector_db=collector_db,
    )
    if not equipments:
        return [], []

    # Flatten results
    eq_results = calc_result if isinstance(calc_result, list) else [calc_result]

    # Calculate contribution map (node_name → contribution)
    equipments_eaf_contribution = calculate_asset_eaf_contributions(
        plant_result=plant_result, eq_results=eq_results
    )

    # Build base result list
    result = []
    for equipment in equipments:
        contribution = equipments_eaf_contribution.get(equipment.location_tag, 0.0)
        total_cost = (equipment.overhaul_cost or 0) + (equipment.service_cost or 0)

        result.append(
            {
                "id": equipment.id,
                "location_tag": equipment.location_tag,
                "name": equipment.equipment_name,
                "total_cost": total_cost,
                "eaf_contribution": contribution,
            }
        )

    # Normalize contribution so sum = 1.0
    total_contrib = sum(item["eaf_contribution"] for item in result) or 1.0
    for item in result:
        item["contribution_norm"] = item["eaf_contribution"] / total_contrib

    # Calculate efficiency and composite score
    for item in result:
        cost = item["total_cost"] or 1.0
        efficiency = item["contribution_norm"] / cost
        item["priority_score"] = item["contribution_norm"] 

    # Choose method
    if use_optimal:
        priority_list, consequence_list = knapsack_selection(result, cost_threshold)
    else:
        priority_list, consequence_list = greedy_selection(result, cost_threshold)

    return priority_list, consequence_list


def calculate_asset_eaf_contributions(plant_result, eq_results):
    """
    Calculate each asset's negative contribution to plant EAF.
    Key assumption: eq_results have aeros_node.node_name matching equipment.location_tag.
    """
    results = defaultdict(float)
    for asset in eq_results:
        node_name = asset.get("aeros_node", {}).get("node_name")
        if node_name:
            results[node_name] = asset.get("contribution_factor", 0.0)
    return results


def greedy_selection(equipments: List[dict], budget: float) -> Tuple[List[dict], List[dict]]:
    """Greedy fallback: select items by score until budget is used."""
    # Sort by priority_score descending
    equipments_sorted = sorted(equipments, key=lambda x: x["priority_score"], reverse=True)
    total_cost = 0
    selected, excluded = [], []
    
    for eq in equipments_sorted:
        if total_cost + eq["cost"] <= budget:
            selected.append(eq)
            total_cost += eq["cost"]
        else:
            excluded.append(eq)
    
    return selected, excluded


def knapsack_selection(equipments: List[dict], budget: float, scale: int = 10_000_000) -> Tuple[List[dict], List[dict]]:
    """
    Select equipment optimally within budget using 0/1 knapsack DP.
    Uses scaling + 1D DP optimization to avoid MemoryError.
    Falls back to greedy if W is too large.
    """
    n = len(equipments)

    # Scale costs + budget
    costs = [int(eq["total_cost"] // scale) for eq in equipments]
    values = [eq["priority_score"] for eq in equipments]
    W = int(budget // scale)

    # Fallback if W is still too large
    if W > 1_000_000:
        print("too big")
        return greedy_selection(equipments, budget)

    # 1D DP array
    dp = [0.0] * (W + 1)
    keep = [[False] * (W + 1) for _ in range(n)]  # track selection choices

    for i in range(n):
        cost, value = costs[i], values[i]
        for w in range(W, cost - 1, -1):
            if dp[w - cost] + value >= dp[w]:  # <= FIXED HERE
                dp[w] = dp[w - cost] + value
                keep[i][w] = True

    # Backtrack to find selected items
    selected, excluded = [], []
    w = W
    for i in range(n - 1, -1, -1):
        if keep[i][w]:
            selected.append(equipments[i])
            w -= costs[i]
        else:
            excluded.append(equipments[i])

    # Optional: fill leftover budget with zero-priority items
    remaining_budget = budget - sum(eq["total_cost"] for eq in selected)
    if remaining_budget > 0:
        for eq in excluded[:]:
            if eq["total_cost"] <= remaining_budget:
                selected.append(eq)
                excluded.remove(eq)
                remaining_budget -= eq["total_cost"]

    return selected, excluded