from __future__ import annotations

from datetime import datetime, timezone

from epochler.contract.models import (
    ContractPhase,
    FrozenContract,
    LockState,
    PhaseArtifact,
    PhaseExecutionStatus,
    PhaseName,
    ensure_builtin_phases,
)


def lock_dataset(contract: FrozenContract) -> FrozenContract:
    ensure_builtin_phases(contract)
    if contract.dataset:
        contract.dataset.lock_state = LockState.FROZEN
        contract.phases[PhaseName.DISCOVERY.value].lock_state = LockState.FROZEN
    return contract


def lock_eval(contract: FrozenContract) -> FrozenContract:
    ensure_builtin_phases(contract)
    if contract.eval_spec:
        contract.eval_spec.lock_state = LockState.FROZEN
        contract.phases[PhaseName.EVALUATION.value].lock_state = LockState.FROZEN
    return contract


def lock_model(contract: FrozenContract) -> FrozenContract:
    ensure_builtin_phases(contract)
    if contract.model_spec:
        contract.model_spec.lock_state = LockState.FROZEN
        contract.phases[PhaseName.TRAINING.value].lock_state = LockState.FROZEN
        contract.phases[PhaseName.ANALYSIS.value].lock_state = LockState.FROZEN
    return contract


def lock_preprocessing(contract: FrozenContract) -> FrozenContract:
    ensure_builtin_phases(contract)
    if contract.preprocessing:
        contract.preprocessing.lock_state = LockState.FROZEN
        contract.phases[PhaseName.PREPROCESSING.value].lock_state = LockState.FROZEN
        contract.phases[PhaseName.FEATURIZATION.value].lock_state = LockState.FROZEN
    return contract


def sync_phase_states(contract: FrozenContract) -> FrozenContract:
    ensure_builtin_phases(contract)

    discovery = contract.phases[PhaseName.DISCOVERY.value]
    if not discovery.invalidated:
        discovery.lock_state = contract.dataset.lock_state if contract.dataset else LockState.HINT
    if discovery.status == PhaseExecutionStatus.PENDING and contract.dataset:
        discovery.status = PhaseExecutionStatus.READY

    preprocessing = contract.phases[PhaseName.PREPROCESSING.value]
    if not preprocessing.invalidated:
        preprocessing.lock_state = contract.preprocessing.lock_state if contract.preprocessing else LockState.HINT

    featurization = contract.phases[PhaseName.FEATURIZATION.value]
    if not featurization.invalidated:
        featurization.lock_state = preprocessing.lock_state

    training = contract.phases[PhaseName.TRAINING.value]
    if not training.invalidated:
        training.lock_state = contract.model_spec.lock_state if contract.model_spec else LockState.HINT

    evaluation = contract.phases[PhaseName.EVALUATION.value]
    if not evaluation.invalidated:
        evaluation.lock_state = contract.eval_spec.lock_state if contract.eval_spec else LockState.HINT

    analysis = contract.phases[PhaseName.ANALYSIS.value]
    if not analysis.invalidated:
        analysis.lock_state = LockState.FROZEN if (
            training.lock_state == LockState.FROZEN and evaluation.lock_state == LockState.FROZEN
        ) else LockState.PROPOSED if (
            training.lock_state == LockState.PROPOSED or evaluation.lock_state == LockState.PROPOSED
        ) else LockState.HINT

    if contract.frozen:
        for phase_name in contract.phase_order:
            phase = contract.phases[phase_name]
            if phase.lock_state != LockState.HINT and not phase.invalidated:
                phase.lock_state = LockState.FROZEN
                if phase.status == PhaseExecutionStatus.PENDING:
                    phase.status = PhaseExecutionStatus.READY
    return contract


def invalidate_phase(
    contract: FrozenContract,
    phase_name: str,
    reason: str = "",
) -> FrozenContract:
    ensure_builtin_phases(contract)
    phase = contract.phases[phase_name]
    phase.status = PhaseExecutionStatus.INVALIDATED
    phase.invalidated = True
    phase.invalidation_reason = reason
    if phase.lock_state == LockState.FROZEN:
        phase.lock_state = LockState.PROPOSED
    phase.artifact_refs = []
    phase.last_run_at = None
    return contract


def invalidate_downstream_phases(
    contract: FrozenContract,
    phase_name: str,
    reason: str = "",
) -> FrozenContract:
    ensure_builtin_phases(contract)
    start = False
    for candidate in contract.phase_order:
        if candidate == phase_name:
            start = True
        if start:
            invalidate_phase(contract, candidate, reason or f"Upstream phase '{phase_name}' changed")
    contract.frozen = False
    return contract


def mark_phase_completed(
    contract: FrozenContract,
    phase_name: str,
    artifacts: list[PhaseArtifact] | None = None,
    script_path: str = "",
) -> FrozenContract:
    ensure_builtin_phases(contract)
    phase = contract.phases[phase_name]
    phase.status = PhaseExecutionStatus.COMPLETED
    phase.invalidated = False
    phase.invalidation_reason = ""
    phase.last_run_at = datetime.now(timezone.utc)
    if artifacts is not None:
        phase.artifact_refs = artifacts
    if script_path:
        phase.script_path = script_path
    return contract


def freeze_contract(contract: FrozenContract) -> FrozenContract:
    """Freeze the entire contract. All components must be at least proposed."""
    sync_phase_states(contract)
    errors: list[str] = []
    if not contract.dataset:
        errors.append("Dataset is not set")
    if not contract.eval_spec:
        errors.append("Eval spec is not set")
    if not contract.model_spec:
        errors.append("Model spec is not set")
    if contract.target_threshold is None:
        errors.append("Target threshold is not set")
    if errors:
        raise ValueError(f"Cannot freeze contract: {'; '.join(errors)}")

    if contract.dataset:
        contract.dataset.lock_state = LockState.FROZEN
    if contract.eval_spec:
        contract.eval_spec.lock_state = LockState.FROZEN
    if contract.model_spec:
        contract.model_spec.lock_state = LockState.FROZEN
    if contract.preprocessing:
        contract.preprocessing.lock_state = LockState.FROZEN
    if contract.reference:
        contract.reference.lock_state = LockState.FROZEN
    for phase_name in contract.phase_order:
        phase = contract.phases[phase_name]
        if phase.lock_state != LockState.HINT:
            phase.lock_state = LockState.FROZEN
            if phase.status == PhaseExecutionStatus.PENDING:
                phase.status = PhaseExecutionStatus.READY
    contract.frozen = True
    return contract


def is_ready_to_run(contract: FrozenContract) -> bool:
    sync_phase_states(contract)
    required = [
        contract.phases[PhaseName.PREPROCESSING.value],
        contract.phases[PhaseName.FEATURIZATION.value],
        contract.phases[PhaseName.TRAINING.value],
        contract.phases[PhaseName.EVALUATION.value],
    ]
    return all(phase.lock_state == LockState.FROZEN for phase in required)


def get_lock_summary(contract: FrozenContract) -> dict[str, str]:
    """Return lock state for each contract component."""
    sync_phase_states(contract)
    return {
        "dataset": contract.dataset.lock_state.value if contract.dataset else "missing",
        "eval_spec": contract.eval_spec.lock_state.value if contract.eval_spec else "missing",
        "model_spec": contract.model_spec.lock_state.value if contract.model_spec else "missing",
        "preprocessing": contract.preprocessing.lock_state.value if contract.preprocessing else "missing",
        "reference": contract.reference.lock_state.value if contract.reference else "n/a",
        "target": "set" if contract.target_threshold is not None else "missing",
        "threshold": "set" if contract.target_threshold is not None else "missing",
        "discovery": contract.phases[PhaseName.DISCOVERY.value].lock_state.value,
        "featurization_or_materialization": contract.phases[PhaseName.FEATURIZATION.value].lock_state.value,
        "training": contract.phases[PhaseName.TRAINING.value].lock_state.value,
        "evaluation": contract.phases[PhaseName.EVALUATION.value].lock_state.value,
        "analysis": contract.phases[PhaseName.ANALYSIS.value].lock_state.value,
        "frozen": str(contract.frozen),
    }
