from __future__ import annotations

import itertools
import time
from dataclasses import dataclass, field
from typing import Any

from app.modules.zootech.lab.calc.engine import calculate_ration
from app.modules.zootech.lab.calc.feed_groups import (
    FEED_GROUPS,
    flatten_group_selections,
    triplet_meets_group_rules,
    validate_group_selections,
)
from app.modules.zootech.lab.calc.feed_groups import classify_feed_group as _classify_feed_group
from app.modules.zootech.lab.calc.formulate_optimize import optimize_shares_for_triplet
from app.modules.zootech.lab.calc.formulate_score import (
    build_calc_lines,
    build_triplet_score_model,
    score_model_shares,
    violation_score,
)
from app.modules.zootech.lab.calc.formulate_validate import validate_component, validate_components
from app.modules.zootech.lab.calc.norms_resolver import NormsParams, NormsResolveRequest, normalize_norms_method, resolve_norms
from app.modules.zootech.lab.calc.nutrients import get_nutrient_value
from app.modules.zootech.lab.indicators import RATION_ALL_INDICATORS
from app.modules.zootech.lab.models import LabAnimalProfile
from app.modules.zootech.lab.services.component_nutrients import nutrients_calc_dict_batch
from app.modules.zootech.lab.services.profile_norms import load_norms_dict
from app.modules.zootech.wesp_bridge_models import Component


@dataclass
class FormulateRequest:
    profile_id: str
    candidate_ids: list[str] = field(default_factory=list)
    group_selections: dict[str, list[str]] = field(default_factory=dict)
    main_feed_ids: list[str] = field(default_factory=list)  # legacy → rough
    mass_kg: float | None = None
    milk_yield_kg: float | None = None
    heads_per_trip: int | None = None
    total_kg_per_head: float = 7.3
    optimize_keys: list[str] = field(default_factory=list)
    objective: str = "min_cost"
    cost_weight: float = 100.0
    grid_step: float = 0.1
    prefilter_k: int = 18
    min_share: float = 0.05
    norms_method: str = "wesp"
    norms_params: dict[str, Any] = field(default_factory=dict)


_DEFAULT_OPTIMIZE_KEYS = (
    "dry_matter",
    "usp",
    "nel",
    "crude_protein",
    "rnb",
    "nfc_pct_dm_uk",
)


def _indicator_def(key: str) -> dict[str, Any] | None:
    for defn in RATION_ALL_INDICATORS:
        if defn["key"] == key:
            return defn
    return None


def _resolve_norms(profile: LabAnimalProfile, req: FormulateRequest) -> tuple[dict, dict]:
    stored = load_norms_dict(profile.id)
    mass = req.mass_kg if req.mass_kg is not None else profile.mass_kg
    milk = req.milk_yield_kg if req.milk_yield_kg is not None else profile.milk_yield_kg
    from app.modules.zootech.lab.services.norms_params import load_norms_params

    method = normalize_norms_method(req.norms_method or profile.norms_method)
    params = load_norms_params(profile)
    if req.norms_params:
        merged = {
            "milkFatPct": params.milk_fat_pct,
            "lactationNo": params.lactation_no,
            "lactationStage": params.lactation_stage,
            "bodyCondition": params.body_condition,
            "housingSystem": params.housing_system,
            "koncOeSv": params.konc_oe_sv,
        }
        merged.update(req.norms_params)
        params = NormsParams.from_dict(merged)
    resolved, meta = resolve_norms(
        NormsResolveRequest(
            method=method,
            stored=stored,
            mass_kg=mass,
            milk_yield_kg=milk,
            ration_type=profile.ration_type,
            force_dynamic=method == "wesp",
            params=params,
        )
    )
    dynamic = meta.get("dynamicNorms") or meta.get("dynamic") or {}
    return resolved, {"normsMethod": method, "dynamicNorms": dynamic, "normsMeta": meta.get("meta")}


def _rough_component_value(
    comp: Component,
    optimize_keys: list[str],
    norms: dict[str, dict[str, float | None]],
    nutrient_cache: dict[str, dict[str, float]],
) -> float | None:
    nutrients = nutrient_cache.get(comp.id, {})
    dm_pct = comp.dry_matter
    total = 0.0
    counted = 0
    for key in optimize_keys:
        defn = _indicator_def(key)
        if defn is None or defn.get("derived"):
            continue
        bounds = norms.get(key) or {}
        target_min = bounds.get("min")
        target_max = bounds.get("max")
        if target_min is None and target_max is None:
            continue
        target = None
        if target_min is not None and target_max is not None:
            target = (float(target_min) + float(target_max)) / 2.0
        elif target_max is not None:
            target = float(target_max) * 0.9
        elif target_min is not None:
            target = float(target_min) * 1.1
        if target is None or target == 0:
            continue
        val = get_nutrient_value(
            dm_pct,
            nutrients,
            defn.get("nutrient_keys") or [],
            indicator_key=key,
        )
        if val is None:
            continue
        rel = (float(val) - target) / abs(target)
        total += rel * rel
        counted += 1
    return total / counted if counted else None


def _prefilter_candidates(
    candidates: list[Component],
    optimize_keys: list[str],
    norms: dict[str, dict[str, float | None]],
    nutrient_cache: dict[str, dict[str, float]],
    *,
    prefilter_k: int,
    must_keep_ids: set[str] | None = None,
) -> list[Component]:
    must_keep_ids = must_keep_ids or set()
    pinned = [c for c in candidates if c.id in must_keep_ids]
    rest = [c for c in candidates if c.id not in must_keep_ids]
    slots = max(prefilter_k - len(pinned), 0)
    if len(candidates) <= prefilter_k:
        return candidates
    if slots <= 0:
        return pinned[:prefilter_k]

    prices = [float(c.price) for c in rest if c.price is not None]
    median_price = sorted(prices)[len(prices) // 2] if prices else 1.0
    if median_price <= 0:
        median_price = 1.0

    scored: list[tuple[float, Component]] = []
    for comp in rest:
        price = float(comp.price) if comp.price is not None else median_price
        cost_part = price / median_price
        nutrient_part = _rough_component_value(comp, optimize_keys, norms, nutrient_cache)
        if nutrient_part is None:
            score = cost_part
        else:
            score = 0.4 * cost_part + 0.6 * nutrient_part
        scored.append((score, comp))
    scored.sort(key=lambda x: x[0])
    return pinned + [comp for _, comp in scored[:slots]]


def _load_profile(profile_id: str) -> LabAnimalProfile:
    profile = LabAnimalProfile.query.filter_by(id=profile_id, is_deleted=False).first()
    if profile is None:
        raise LookupError("Профиль не найден")
    return profile


def _load_eligible_components(candidate_ids: list[str]) -> tuple[list[Component], list[dict[str, Any]]]:
    unique = list(dict.fromkeys(candidate_ids))
    validations = validate_components(unique)
    ineligible = [v for v in validations if not v["eligible"]]
    if ineligible:
        raise ValueError("ineligible_components", ineligible)
    comps: list[Component] = []
    for cid in unique:
        comp = Component.query.filter_by(id=cid, is_deleted=False).first()
        if comp is None:
            raise ValueError("component_not_found", cid)
        comps.append(comp)
    return comps, validations


def _resolve_group_selections(req: FormulateRequest) -> dict[str, list[str]]:
    selections = dict(req.group_selections or {})
    if req.main_feed_ids and not selections.get("rough"):
        selections["rough"] = list(req.main_feed_ids)
    return selections


def formulate(req: FormulateRequest) -> dict[str, Any]:
    group_selections = _resolve_group_selections(req)
    candidate_ids = list(req.candidate_ids)
    if group_selections:
        pool_errors = validate_group_selections(group_selections)
        if pool_errors:
            raise ValueError("group_selection_invalid", pool_errors)
        candidate_ids = flatten_group_selections(group_selections)

    if len(candidate_ids) < 3:
        raise ValueError("candidate_ids_min_3")
    if len(set(candidate_ids)) != len(candidate_ids):
        raise ValueError("candidate_ids_duplicate")

    profile = _load_profile(req.profile_id)
    candidates, _ = _load_eligible_components(candidate_ids)
    nutrient_cache = nutrients_calc_dict_batch(candidate_ids)
    norms, dynamic = _resolve_norms(profile, req)
    optimize_keys = req.optimize_keys or list(_DEFAULT_OPTIMIZE_KEYS)
    heads = max(int(req.heads_per_trip or 10), 1)
    herd_scale = float(req.total_kg_per_head) * heads
    profile_mass_kg = req.mass_kg if req.mass_kg is not None else profile.mass_kg

    must_keep: set[str] = set()
    for g in FEED_GROUPS:
        if g["required"]:
            must_keep.update(group_selections.get(g["id"]) or [])

    prefilter_applied = len(candidates) > req.prefilter_k
    shortlist = _prefilter_candidates(
        candidates,
        optimize_keys,
        norms,
        nutrient_cache,
        prefilter_k=req.prefilter_k,
        must_keep_ids=must_keep,
    )

    started = time.perf_counter()
    evaluations = 0
    triplets_evaluated = 0
    triplet_winners: list[dict[str, Any]] = []

    def _eval_triplet(
        triplet: tuple[Component, Component, Component],
    ) -> dict[str, Any] | None:
        nonlocal evaluations
        model = build_triplet_score_model(
            triplet,
            nutrient_cache=nutrient_cache,
            norms=norms,
            optimize_keys=optimize_keys,
            herd_scale=herd_scale,
            heads=heads,
            cost_weight=req.cost_weight,
            profile_mass_kg=profile_mass_kg,
        )

        def score_fn(shares: tuple[float, float, float]) -> float:
            _violation, _cost_head, score = score_model_shares(model, shares)
            return score

        opt = optimize_shares_for_triplet(
            triplet,
            score_fn,
            min_share=req.min_share,
            grid_step=req.grid_step,
        )
        if opt is None:
            return None
        evaluations += opt.evaluations
        violation, cost_head, score = score_model_shares(model, opt.shares)
        lines = build_calc_lines(triplet, opt.shares, herd_scale, nutrient_cache)
        daily = [opt.shares[i] * herd_scale for i in range(3)]
        return {
            "score": score,
            "violation": violation,
            "costHead": cost_head,
            "costTotal": cost_head * heads,
            "triplet": triplet,
            "shares": opt.shares,
            "daily": daily,
            "lines": lines,
        }

    for triplet in itertools.combinations(shortlist, 3):
        triplet_ids = {c.id for c in triplet}
        if not triplet_meets_group_rules(triplet_ids, group_selections):
            continue
        triplets_evaluated += 1
        best_triplet_result = _eval_triplet(triplet)
        if best_triplet_result is None:
            continue
        triplet_winners.append(best_triplet_result)

    if not triplet_winners:
        if group_selections:
            raise ValueError("no_feasible_solution_groups")
        raise ValueError("no_feasible_solution")

    triplet_winners.sort(key=lambda x: x["score"])
    best = triplet_winners[0]

    full_calc = calculate_ration(
        profile.ration_type or "DAIRY",
        best["lines"],
        norms,
        heads_per_trip=heads,
        profile_mass_kg=profile_mass_kg,
    )
    best["calc"] = full_calc
    best["violation"] = violation_score(
        full_calc.get("indicators") or [],
        optimize_keys,
        norms,
    )
    best["score"] = best["violation"] * req.cost_weight + best["costHead"]

    duration_ms = int((time.perf_counter() - started) * 1000)

    alternatives = [
        {
            "componentIds": [c.id for c in item["triplet"]],
            "names": [c.name for c in item["triplet"]],
            "score": item["score"],
            "violation": item["violation"],
            "costPerHead": item["costHead"],
        }
        for item in triplet_winners[:3]
    ]
    alternatives.sort(key=lambda x: x["score"])
    alternatives = alternatives[:3]

    result_lines = []
    for i, comp in enumerate(best["triplet"]):
        s1, s2, s3 = best["shares"]
        share = (s1, s2, s3)[i]
        result_lines.append(
            {
                "componentId": comp.id,
                "name": comp.name,
                "dailyKg": round(best["daily"][i], 4),
                "sharePct": round(share * 100, 2),
                "pricePerKg": comp.price,
                "dryMatterPct": comp.dry_matter,
            }
        )

    violation = best["violation"]
    return {
        "lines": result_lines,
        "candidatePoolSize": len(candidates),
        "groupSelections": group_selections,
        "shortlistedIds": [c.id for c in shortlist],
        "costTotal": best["costTotal"],
        "costPerHead": best["costHead"],
        "score": best["score"],
        "violation": violation,
        "feasible": violation < 0.01,
        "indicators": best["calc"].get("indicators") or [],
        "totals": best["calc"].get("totals") or [],
        "optimizeKeys": optimize_keys,
        "dynamicNorms": dynamic.get("dynamicNorms") or None,
        "normsMethod": dynamic.get("normsMethod"),
        "normsMeta": dynamic.get("normsMeta"),
        "alternatives": alternatives,
        "searchStats": {
            "evaluations": evaluations,
            "tripletsEvaluated": triplets_evaluated,
            "durationMs": duration_ms,
            "prefilterApplied": prefilter_applied,
            "prefilterK": req.prefilter_k,
            "scoreEngine": "fast",
            "optimizer": "slsqp",
            "normsMethod": dynamic.get("normsMethod"),
        },
    }


def list_formulate_components() -> list[dict[str, Any]]:
    rows = (
        Component.query.filter_by(is_active=True, is_deleted=False)
        .order_by(Component.name)
        .all()
    )
    out: list[dict[str, Any]] = []
    for comp in rows:
        v = validate_component(comp.id)
        feed_group = _classify_feed_group(comp)
        out.append(
            {
                "id": comp.id,
                "name": comp.name,
                "type": comp.type,
                "feedGroup": feed_group,
                "eligible": v["eligible"],
                "missing": v["missing"],
                "warnings": v["warnings"],
                "dryMatterPct": v["dryMatterPct"],
                "hasPrice": v["hasPrice"],
                "price": v["price"],
                "mainFeedDmGPerKg": v.get("mainFeedDmGPerKg"),
                "isMainFeed": v.get("isMainFeed", False),
            }
        )
    return out