from __future__ import annotations

import uuid
from typing import Any

from app.modules.zootech.wesp_bridge_db import db
from app.modules.zootech.lab.calc.feed_groups import classify_feed_group
from app.modules.zootech.lab.calc.gfe_policies import DeriveContext
from app.modules.zootech.lab.calc.ingredient_catalog import DERIVED_HEADERS
from app.modules.zootech.lab.calc.ingredient_derive import derive_ingredient_nutrients
from app.modules.zootech.lab.models import LabComponentNutrientValue
from app.modules.zootech.lab.nutrient_keys import augment_with_indicator_keys, canonicalize_for_storage
from app.modules.zootech.lab.nutrient_schema import (
    dry_matter_g_per_kg,
    mapping_to_calc_dict,
    read_from_mapping,
)
from app.modules.zootech.wesp_bridge_models import Component

_DERIVE_INPUT_KEYS = ("СВ", "Сыр. Протеин", "Сырая клетч", "Сырой жир")
_EAV_SKIP_KEYS = frozenset({"СВ"})
_MAIN_FEED_KEYS = ("Осн.Корм", "СВ Основной корм")


def derive_context_for_component(
    component_id: str | None,
    merged: dict[str, Any] | None = None,
) -> DeriveContext:
    """Контекст derive: тип корма + признак основного корма."""
    main_feed = read_from_mapping(merged or {}, _MAIN_FEED_KEYS)
    is_main = main_feed is not None and float(main_feed) > 0
    feed_group: str = "unknown"
    if component_id:
        comp = Component.query.filter_by(id=component_id, is_deleted=False).first()
        if comp is not None:
            feed_group = classify_feed_group(comp)
    return DeriveContext(feed_group=feed_group, is_main_feed=is_main)  # type: ignore[arg-type]


def _component_dry_matter_pct(component_id: str | None) -> float | None:
    if not component_id:
        return None
    comp = Component.query.get(component_id)
    return comp.dry_matter if comp else None


def _inject_sv_for_derive(merged: dict[str, Any], dry_matter_pct: float | None) -> dict[str, Any]:
    out = dict(merged)
    sv = dry_matter_g_per_kg(dry_matter_pct)
    if sv is not None:
        out["СВ"] = sv
    return out


def _should_derive(merged: dict[str, Any]) -> bool:
    """Derive при полном базовом вводе (СВ из component.dry_matter + 3 показателя)."""
    return all(read_from_mapping(merged, (k,)) is not None for k in _DERIVE_INPUT_KEYS)


def nutrients_full_dict(component_id: str | None) -> dict[str, float]:
    if not component_id:
        return {}
    rows = LabComponentNutrientValue.query.filter_by(component_id=component_id).all()
    return {
        row.nutrient_key: float(row.value)
        for row in rows
        if row.value is not None and row.nutrient_key not in _EAV_SKIP_KEYS
    }


def nutrients_api_dict(component_id: str | None) -> dict[str, float]:
    return nutrients_full_dict(component_id)


def _fill_missing_derived(
    merged: dict[str, float],
    *,
    component_id: str | None,
) -> dict[str, float]:
    """Добавляет derived-поля (пОВ, переваримые фракции), не перезаписывая введённые лаб. значения."""
    if not _should_derive(merged):
        return merged
    ctx = derive_context_for_component(component_id, merged)
    derived = derive_ingredient_nutrients(merged, context=ctx)
    out = dict(merged)
    for header in DERIVED_HEADERS:
        if read_from_mapping(out, (header,)) is not None:
            continue
        val = read_from_mapping(derived, (header,))
        if val is not None:
            out[header] = float(val)
    return out


def _repair_for_ration_calc(
    full: dict[str, float],
    dry_matter_pct: float | None,
    component_id: str | None = None,
) -> dict[str, float]:
    merged = _inject_sv_for_derive(full, dry_matter_pct)
    oe = read_from_mapping(merged, ("ОЭ-КРС", "    ОЭ-КРС"))
    nel = read_from_mapping(merged, ("ЧЭЛ- КРС", "    ЧЭЛ- КРС"))
    energy_bad = (oe is not None and oe < 0) or (nel is not None and nel < 0)
    if energy_bad:
        minimal = {key: read_from_mapping(merged, (key,)) for key in _DERIVE_INPUT_KEYS}
        if all(v is not None for v in minimal.values()):
            ctx = derive_context_for_component(component_id, merged)
            merged.update(derive_ingredient_nutrients(minimal, context=ctx))
    else:
        merged = _fill_missing_derived(merged, component_id=component_id)
    return merged


def _calc_dict_from_full(
    full: dict[str, float],
    dry_matter_pct: float | None,
    component_id: str | None,
) -> dict[str, float]:
    repaired = _repair_for_ration_calc(full, dry_matter_pct, component_id)
    return mapping_to_calc_dict(augment_with_indicator_keys(repaired))


def nutrients_calc_dict(component_id: str | None) -> dict[str, float]:
    full = nutrients_full_dict(component_id)
    dry_matter_pct = _component_dry_matter_pct(component_id)
    return _calc_dict_from_full(full, dry_matter_pct, component_id)


def nutrients_calc_dict_batch(component_ids: list[str]) -> dict[str, dict[str, float]]:
    """Batch-load EAV + dry_matter for formulate (one SQL round-trip per table)."""
    unique = list(dict.fromkeys(cid for cid in component_ids if cid))
    if not unique:
        return {}

    eav_by_id: dict[str, dict[str, float]] = {cid: {} for cid in unique}
    rows = LabComponentNutrientValue.query.filter(
        LabComponentNutrientValue.component_id.in_(unique),
    ).all()
    for row in rows:
        if row.value is None or row.nutrient_key in _EAV_SKIP_KEYS:
            continue
        eav_by_id.setdefault(row.component_id, {})[row.nutrient_key] = float(row.value)

    dm_by_id: dict[str, float | None] = {cid: None for cid in unique}
    for comp in Component.query.filter(
        Component.id.in_(unique),
        Component.is_deleted.is_(False),
    ).all():
        dm_by_id[comp.id] = comp.dry_matter

    return {
        cid: _calc_dict_from_full(eav_by_id.get(cid, {}), dm_by_id.get(cid), cid)
        for cid in unique
    }


def nutrients_is_empty(component_id: str | None) -> bool:
    if not component_id:
        return True
    return (
        LabComponentNutrientValue.query.filter(
            LabComponentNutrientValue.component_id == component_id,
            LabComponentNutrientValue.nutrient_key.notin_(_EAV_SKIP_KEYS),
        ).first()
        is None
    )


def _upsert_eav(component_id: str, nutrients: dict[str, float]) -> None:
    existing = {
        row.nutrient_key: row
        for row in LabComponentNutrientValue.query.filter_by(component_id=component_id).all()
    }
    for key in list(existing):
        if key in _EAV_SKIP_KEYS:
            db.session.delete(existing[key])
    existing = {k: v for k, v in existing.items() if k not in _EAV_SKIP_KEYS}
    for key, value in nutrients.items():
        if value is None or key in _EAV_SKIP_KEYS:
            continue
        row = existing.get(key)
        if row is None:
            row = LabComponentNutrientValue(
                id=str(uuid.uuid4()),
                component_id=component_id,
                nutrient_key=key,
                value=float(value),
            )
            db.session.add(row)
            existing[key] = row
        else:
            row.value = float(value)


def _prepare_payload(
    component_id: str,
    nutrients: dict[str, Any] | None,
    *,
    pin: bool = False,
) -> dict[str, float]:
    dry_matter_pct = _component_dry_matter_pct(component_id)
    merged = _inject_sv_for_derive(
        {**nutrients_full_dict(component_id), **(nutrients or {})},
        dry_matter_pct,
    )
    if pin:
        payload = merged
    elif _should_derive(merged):
        ctx = derive_context_for_component(component_id, merged)
        derived = derive_ingredient_nutrients(merged, context=ctx)
        payload = {**merged, **derived}
    else:
        payload = merged
    stored = canonicalize_for_storage(payload)
    return {k: v for k, v in stored.items() if k not in _EAV_SKIP_KEYS}


def save_component_nutrients(
    component_id: str,
    nutrients: dict[str, Any] | None,
    *,
    user_id: str = "system",
    pin: bool = False,
) -> dict[str, Any]:
    """Единственная точка записи EAV + derive. Возвращает stored dict и affectedRecipeIds."""
    stored = _prepare_payload(component_id, nutrients, pin=pin)
    _upsert_eav(component_id, stored)
    from app.modules.zootech.lab.services.ration_recalc import find_rations_by_component

    affected = find_rations_by_component(component_id)
    return {"stored": stored, "affectedRecipeIds": affected, "userId": user_id}


# Backward-compatible aliases for tests and gradual migration
def upsert_from_api_dict(component_id: str, nutrients: dict[str, Any] | None) -> dict[str, float]:
    result = save_component_nutrients(component_id, nutrients)
    return result["stored"]


def pin_nutrient_values(component_id: str, values: dict[str, float]) -> None:
    save_component_nutrients(component_id, values, pin=True)


def upsert_from_column_values(component_id: str, values: dict[str, float | None]) -> dict[str, float]:
    from app.modules.zootech.lab.nutrient_schema import column_values_to_mapping

    merged = nutrients_full_dict(component_id)
    merged.update(column_values_to_mapping(values))
    return upsert_from_api_dict(component_id, merged)