from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Any, Mapping

from .knob_descriptor import KnobDescriptor


@dataclass(frozen=True)
class CoordinateSearchPolicy:
    initial_relative_step: float = 1.0
    initial_fraction_step: float = 0.05
    step_multipliers: tuple[float, ...] = (1.0,)
    grow_factor: float = 1.5
    shrink_factor: float = 0.5
    min_score: float = 0.0


@dataclass(frozen=True)
class CoordinateOperatorState:
    knob: str
    direction: str
    trust_radius: float | None = None
    last_good_value: Any | None = None
    last_bad_value: Any | None = None
    tested_values: tuple[Any, ...] = ()


@dataclass(frozen=True)
class MechanismCandidate:
    action_id: str
    knob: str
    mechanism: str
    operator: str
    direction: str
    patch: dict[str, Any]
    score: float
    score_factors: dict[str, float]
    evidence_refs: tuple[str, ...] = ()


def coordinate_line_search_candidates(
    *,
    current_config: Mapping[str, Any],
    descriptors: tuple[KnobDescriptor, ...],
    evidence_weights: Mapping[str, float],
    states: Mapping[tuple[str, str], CoordinateOperatorState] | None = None,
    policy: CoordinateSearchPolicy | None = None,
) -> tuple[MechanismCandidate, ...]:
    policy = policy or CoordinateSearchPolicy()
    states = states or {}
    candidates: list[MechanismCandidate] = []
    for descriptor in descriptors:
        if "coordinate_line_search" not in descriptor.operators:
            continue
        direction, mechanism, evidence = _choose_direction(descriptor, evidence_weights)
        if direction is None or mechanism is None:
            continue
        if evidence < policy.min_score:
            continue
        state = states.get((descriptor.name, direction))
        current = descriptor.current_value(current_config)
        multipliers = (
            (1.0,) if descriptor.search_geometry == "toggle" else policy.step_multipliers
        )
        for multiplier in multipliers:
            target = _propose_value(
                descriptor=descriptor,
                current=current,
                direction=direction,
                state=state,
                policy=policy,
                step_multiplier=multiplier,
            )
            if target is None or target == current:
                continue
            risk = _direction_risk(descriptor, direction, evidence_weights)
            step_risk = _step_risk(multiplier)
            score = max(0.0, evidence - risk - step_risk)
            candidates.append(
                MechanismCandidate(
                    action_id=(
                        f"coordinate_line_search:{descriptor.search_geometry}:"
                        f"{descriptor.name}:{direction}:step={_stable_token(multiplier)}:"
                        f"{_stable_token(current)}->{_stable_token(target)}"
                    ),
                    knob=descriptor.name,
                    mechanism=mechanism,
                    operator="coordinate_line_search",
                    direction=direction,
                    patch={descriptor.name: target},
                    score=round(score, 4),
                    score_factors={
                        "mechanism_evidence": round(evidence, 4),
                        "direction_risk": round(risk, 4),
                        "step_risk": round(step_risk, 4),
                    },
                    evidence_refs=(mechanism,),
                )
            )
    candidates.sort(key=lambda item: (item.score, item.action_id), reverse=True)
    return tuple(candidates)


def _choose_direction(
    descriptor: KnobDescriptor,
    evidence_weights: Mapping[str, float],
) -> tuple[str | None, str | None, float]:
    best_direction: str | None = None
    best_mechanism: str | None = None
    best_weight = 0.0
    for direction, mechanisms in descriptor.directional_effects.items():
        for mechanism in mechanisms:
            weight = float(evidence_weights.get(mechanism, 0.0))
            if weight > best_weight:
                best_direction = direction
                best_mechanism = mechanism
                best_weight = weight
    return best_direction, best_mechanism, best_weight


def _direction_risk(
    descriptor: KnobDescriptor,
    direction: str,
    evidence_weights: Mapping[str, float],
) -> float:
    risks = descriptor.risk_effects.get(direction, ())
    if not risks:
        return 0.0
    return min(0.5, 0.2 * max(float(evidence_weights.get(item, 0.0)) for item in risks))


def _step_risk(multiplier: float) -> float:
    return min(0.2, max(0.0, float(multiplier) - 1.0) * 0.04)


def _propose_value(
    *,
    descriptor: KnobDescriptor,
    current: Any,
    direction: str,
    state: CoordinateOperatorState | None,
    policy: CoordinateSearchPolicy,
    step_multiplier: float,
) -> Any | None:
    if descriptor.search_geometry == "toggle":
        if not isinstance(current, bool):
            current = bool(current)
        return not current
    if descriptor.search_geometry == "bounded_fraction":
        value = _as_float(current)
        if value is None:
            return None
        radius = (
            state.trust_radius
            if state is not None and state.trust_radius is not None
            else policy.initial_fraction_step
        )
        radius *= max(float(step_multiplier), 0.0)
        if direction == "decrease":
            target = value - radius
        else:
            target = value + radius
        return _canonicalize_value(descriptor, target)
    if descriptor.search_geometry == "linear":
        value = _as_float(current)
        if value is None:
            return None
        radius = (
            state.trust_radius
            if state is not None and state.trust_radius is not None
            else 1.0
        )
        radius *= max(float(step_multiplier), 0.0)
        return _canonicalize_value(
            descriptor,
            value - radius if direction == "decrease" else value + radius,
        )
    if descriptor.search_geometry == "positive_capacity":
        value = _as_float(current)
        if value is None or value <= 0:
            min_value = descriptor.constraints.min_value
            return _canonicalize_value(descriptor, min_value if min_value is not None else 1)
        radius = (
            state.trust_radius
            if state is not None and state.trust_radius is not None
            else policy.initial_relative_step
        )
        radius *= max(float(step_multiplier), 0.0)
        factor = max(1.0 + radius, 1.01)
        target = value / factor if direction == "decrease" else value * factor
        return _canonicalize_value(descriptor, target)
    raise ValueError(f"unsupported search geometry {descriptor.search_geometry!r}")


def _canonicalize_value(descriptor: KnobDescriptor, value: float | int) -> Any:
    target = float(value)
    if descriptor.constraints.min_value is not None:
        target = max(target, float(descriptor.constraints.min_value))
    if descriptor.constraints.max_value is not None:
        target = min(target, float(descriptor.constraints.max_value))
    if descriptor.constraints.integer or descriptor.value_type == "int":
        integer_target = int(math.ceil(target))
        multiple_of = descriptor.constraints.multiple_of
        if multiple_of is not None and multiple_of > 1:
            integer_target = int(math.ceil(integer_target / multiple_of) * multiple_of)
        if descriptor.constraints.min_value is not None:
            integer_target = max(integer_target, int(descriptor.constraints.min_value))
        if descriptor.constraints.max_value is not None:
            integer_target = min(integer_target, int(descriptor.constraints.max_value))
        return integer_target
    return round(target, 6)


def _as_float(value: Any) -> float | None:
    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str) and value.strip():
        try:
            return float(value.strip())
        except ValueError:
            return None
    return None


def _stable_token(value: Any) -> str:
    return repr(value)
