
from app.api.carpentry import schema
from app.api.commercial_equipment.schemas import CommercialEquipmentCreate
from app.api.network.schemas import NetworkPlanningCreate
from app.models.main.catchment_potential import CatchmentPotentialBase, TblCatchmentPotential
from app.models.main.customer_location import CustomerLocationBase, TblCustomerLocation
import secrets
import string

def calculate_total_household_consumption(data: CustomerLocationBase) -> int:
    if data.shopping_frequency_per_month is None or data.household_consumption_per_month is None:
        raise ValueError("shopping_frequency_per_month and household_consumption_per_month must not be None")
    return data.shopping_frequency_per_month * data.household_consumption_per_month

def fill_and_calculate_total_household_consumption(data: dict, existing: "TblCustomerLocation") -> dict:
    """
    Merge partial update data with existing DB values and calculate total_household_consumption_per_month.
    """
    shopping_freq = data.get("shopping_frequency_per_month") or existing.shopping_frequency_per_month
    household_cons = data.get("household_consumption_per_month") or existing.household_consumption_per_month

    if shopping_freq is None or household_cons is None:
        raise ValueError("shopping_frequency_per_month and household_consumption_per_month must not be None")

    data["total_household_consumption_per_month"] = shopping_freq * household_cons
    return data

def calculate_potential_number_of_household_consumption(data: CatchmentPotentialBase) -> int:
    if data.number_of_households is None or data.percentage_of_segment is None:
        raise ValueError("number_of_households and percentage_of_segment must not be None")
    return int(data.number_of_households * data.percentage_of_segment)

def fill_and_calculate_potential_households(data: dict, existing: TblCatchmentPotential) -> dict:
    """
    Merge partial update data with existing DB values and calculate potential_number_of_households.
    """
    number_of_households = data.get("number_of_households") or existing.number_of_households
    percentage_of_segment = data.get("percentage_of_segment") or existing.percentage_of_segment

    if number_of_households is None or percentage_of_segment is None:
        raise ValueError("number_of_households and percentage_of_segment must not be None")

    data["potential_number_of_households"] = int(number_of_households * percentage_of_segment)
    return data

def calculate_total_store_counts(data: NetworkPlanningCreate) -> dict:
    if any(
        getattr(data, field) is None
        for field in [
            "format_a_y1", "format_a_y2", "format_a_y3", "format_a_y4", "format_a_y5",
            "format_b_y1", "format_b_y2", "format_b_y3", "format_b_y4", "format_b_y5"
        ]
    ):
        raise ValueError("All format_a and format_b fields for Y1–Y5 must be provided")

    return {
        "total_store_count_y1": data.format_a_y1 + data.format_b_y1,
        "total_store_count_y2": data.format_a_y2 + data.format_b_y2,
        "total_store_count_y3": data.format_a_y3 + data.format_b_y3,
        "total_store_count_y4": data.format_a_y4 + data.format_b_y4,
        "total_store_count_y5": data.format_a_y5 + data.format_b_y5,
    }

def calculate_total_store_counts_update(data: dict, existing: object) -> dict:
    totals = {}

    for year in range(1, 6):
        fa_key = f"format_a_y{year}"
        fb_key = f"format_b_y{year}"
        ts_key = f"total_store_count_y{year}"

        # Get from data (request) or fallback to existing DB value
        fa_val = data.get(fa_key, getattr(existing, fa_key, None))
        fb_val = data.get(fb_key, getattr(existing, fb_key, None))

        if fa_val is not None and fb_val is not None:
            totals[ts_key] = fa_val + fb_val

    return totals

def generate_secure_password(length=12):
    alphabet = string.ascii_letters + string.digits
    return ''.join(secrets.choice(alphabet) for _ in range(length))

def prepare_furniture_data(request: schema.CarpentryCreate) -> dict:
    data = request.model_dump()
    data["total"] = data["cost_per_unit"] * data["number_of_units"]
    return data

def prepare_equipment_data(request: CommercialEquipmentCreate) -> dict:
    data = request.model_dump()
    data["total"] = data["cost_per_unit"] * data["number_of_units"]
    return data
