from __future__ import annotations

from collections import defaultdict
from datetime import timedelta

from .config import GeneratorConfig
from .models import DatasetBundle, stable_id


def generate_mrp(bundle: DatasetBundle, config: GeneratorConfig) -> DatasetBundle:
    materials = bundle.rows("materials")
    work_packages = bundle.rows("work-packages")
    bom_lines = bundle.rows("pbom") or bundle.rows("ebom")
    if not materials or not work_packages or not bom_lines:
        raise ValueError("MRP requires materials, work packages, and BOM rows")
    inventory_by_material: dict[str, float] = defaultdict(float)
    for row in bundle.rows("inventory"):
        if row.get("qualityStatus") == "RELEASED":
            inventory_by_material[str(row["materialId"])] += float(row["quantity"])
    held_allocations: dict[str, float] = defaultdict(float)
    released_allocations: dict[str, float] = defaultdict(float)
    inventory_index = {str(row["inventoryId"]): row for row in bundle.rows("inventory")}
    for allocation in bundle.rows("inventory-allocations"):
        inventory = inventory_index.get(str(allocation["inventoryId"]))
        if not inventory:
            continue
        target = released_allocations if allocation.get("releaseAllowed") else held_allocations
        target[str(inventory["materialId"])] += float(allocation["quantity"])
    receipt_pool: dict[str, float] = defaultdict(float)
    for row in bundle.rows("planned-receipts"):
        if row.get("trusted"):
            receipt_pool[str(row["materialId"])] += float(row["quantity"])
    usable_pool = {mid: max(0.0, qty - held_allocations[mid]) for mid, qty in inventory_by_material.items()}
    material_index = {str(row["materialId"]): row for row in materials}
    bom_by_mode: dict[str, list[dict]] = defaultdict(list)
    for bom in bom_lines:
        material = material_index.get(str(bom["materialId"]))
        if material:
            bom_by_mode[str(material["sourcingMode"])].append(bom)
    requirement_target = config.profile.wbs_task_count
    requirements: list[dict] = []
    package_requirements: dict[str, list[dict]] = defaultdict(list)
    for index in range(requirement_target):
        package = work_packages[index % len(work_packages)]
        if index < config.profile.purchase_suggestion_count and bom_by_mode["BUY"]:
            bom = bom_by_mode["BUY"][index % len(bom_by_mode["BUY"])]
            forced_shortage = True
        elif (
            index < config.profile.purchase_suggestion_count + config.profile.outsource_suggestion_count
            and bom_by_mode["OUTSOURCE"]
        ):
            bom = bom_by_mode["OUTSOURCE"][index % len(bom_by_mode["OUTSOURCE"])]
            forced_shortage = True
        else:
            bom = bom_lines[index % len(bom_lines)]
            forced_shortage = False
        material_id = str(bom["materialId"])
        material = material_index[material_id]
        gross = round(max(1.0, float(bom.get("quantity") or 1.0)) * (1 + index % 4), 3)
        if forced_shortage:
            gross = round(gross + 250.0 + index % 50, 3)
        safety = float(material.get("safetyStock") or 0)
        scrap = round(gross * float(material.get("scrapRate") or 0), 3)
        total = round(gross + safety + scrap, 3)
        stock_used = round(min(usable_pool.get(material_id, 0.0), total), 3)
        usable_pool[material_id] = round(usable_pool.get(material_id, 0.0) - stock_used, 3)
        remaining = round(total - stock_used, 3)
        allocation_used = round(min(released_allocations.get(material_id, 0.0), remaining), 3)
        released_allocations[material_id] = round(released_allocations.get(material_id, 0.0) - allocation_used, 3)
        remaining = round(remaining - allocation_used, 3)
        receipt_used = round(min(receipt_pool.get(material_id, 0.0), remaining), 3)
        receipt_pool[material_id] = round(receipt_pool.get(material_id, 0.0) - receipt_used, 3)
        remaining = round(remaining - receipt_used, 3)
        substitute_used = 0.0
        net = max(0.0, round(remaining, 3))
        need_date = package.get("needDate") or (config.planning_base_date + timedelta(days=30 + index % 300)).isoformat()
        row = {
            "requirementId": stable_id("requirement", package["workPackageId"], material_id, index, prefix="REQ"),
            "projectId": package["projectId"],
            "wbsId": package["wbsId"],
            "workPackageId": package["workPackageId"],
            "bomLineId": bom["bomLineId"],
            "materialId": material_id,
            "grossRequirement": gross,
            "safetyStockRequirement": safety,
            "scrapRequirement": scrap,
            "stockUsed": stock_used,
            "releasedAllocationUsed": allocation_used,
            "plannedReceiptUsed": receipt_used,
            "substituteUsed": substitute_used,
            "netRequirement": net,
            "needDate": need_date,
            "sourcingMode": material["sourcingMode"],
            "fulfillmentMode": material["fulfillmentMode"],
            "peggingRef": f"wbs:{package['wbsId']}|bom:{bom['bomLineId']}",
        }
        requirements.append(row)
        package_requirements[str(package["workPackageId"])].append(row)
    readiness: list[dict] = []
    for package in work_packages:
        rows = package_requirements.get(str(package["workPackageId"]), [])
        shortages = sum(1 for row in rows if float(row["netRequirement"]) > 0)
        total = len(rows)
        percent = 100.0 if total == 0 else round(100 * (total - shortages) / total, 2)
        readiness.append({
            "kitReadinessId": stable_id("kit", package["workPackageId"], prefix="KIT"),
            "workPackageId": package["workPackageId"],
            "readinessPercent": percent,
            "readyDate": package.get("needDate"),
            "shortageCount": shortages,
            "status": "READY" if shortages == 0 else "SHORTAGE",
        })
    bundle.set_rows("material-requirements", requirements)
    bundle.set_rows("kit-readiness", readiness)
    return bundle