from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from epochler.contract.engine import invalidate_downstream_phases, sync_phase_states
from epochler.contract.models import (
    Conversation,
    ConversationState,
    DatasetAsset,
    DecisionRecord,
    DecisionStatus,
    DecisionType,
    EvalSpec,
    LockState,
    ModelSpec,
    NextAction,
    NextActionType,
    PhaseExecutionStatus,
    PhaseName,
    PreprocessingPlan,
    Readiness,
)


_DECISION_PHASES: dict[DecisionType, list[str]] = {
    DecisionType.DATASET: [PhaseName.DISCOVERY.value],
    DecisionType.PREPROCESSING: [
        PhaseName.PREPROCESSING.value,
        PhaseName.FEATURIZATION.value,
    ],
    DecisionType.MODEL: [PhaseName.TRAINING.value, PhaseName.ANALYSIS.value],
    DecisionType.EVALUATION: [PhaseName.EVALUATION.value, PhaseName.ANALYSIS.value],
    DecisionType.TARGET: [PhaseName.ANALYSIS.value],
    DecisionType.RUN_PLAN: [],
    DecisionType.POST_RUN_RECOMMENDATION: [PhaseName.ANALYSIS.value],
    DecisionType.GENERAL: [],
}


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _decision_snapshot(contract, decision_type: DecisionType) -> dict[str, Any]:
    if decision_type == DecisionType.DATASET:
        return contract.dataset.model_dump(mode="json") if contract.dataset else {}
    if decision_type == DecisionType.PREPROCESSING:
        return contract.preprocessing.model_dump(mode="json") if contract.preprocessing else {}
    if decision_type == DecisionType.MODEL:
        return contract.model_spec.model_dump(mode="json") if contract.model_spec else {}
    if decision_type == DecisionType.EVALUATION:
        return contract.eval_spec.model_dump(mode="json") if contract.eval_spec else {}
    if decision_type == DecisionType.TARGET:
        return {
            "target_threshold": contract.target_threshold,
            "threshold_direction": contract.threshold_direction,
            "budget_minutes": contract.budget_minutes,
            "max_experiments": contract.max_experiments,
        }
    return {}


def _set_component_state(contract, decision_type: DecisionType, lock_state: LockState) -> None:
    if decision_type == DecisionType.DATASET and contract.dataset:
        contract.dataset.lock_state = lock_state
    elif decision_type == DecisionType.PREPROCESSING and contract.preprocessing:
        contract.preprocessing.lock_state = lock_state
    elif decision_type == DecisionType.MODEL and contract.model_spec:
        contract.model_spec.lock_state = lock_state
    elif decision_type == DecisionType.EVALUATION and contract.eval_spec:
        contract.eval_spec.lock_state = lock_state


def _invalidate_for_decision(conversation: Conversation, decision_type: DecisionType, reason: str) -> None:
    for phase_name in _DECISION_PHASES.get(decision_type, []):
        invalidate_downstream_phases(conversation.contract, phase_name, reason)


def add_or_replace_decision(
    conversation: Conversation,
    decision: DecisionRecord,
) -> DecisionRecord:
    decision.after = decision.after or _decision_snapshot(conversation.contract, decision.type)
    decision.phase_impacts = decision.phase_impacts or _DECISION_PHASES.get(decision.type, [])
    decision.updated_at = _now()
    for existing in conversation.decisions:
        if existing.type == decision.type and existing.status in {
            DecisionStatus.PROPOSED,
            DecisionStatus.NEEDS_INPUT,
            DecisionStatus.BLOCKED,
        }:
            existing.status = DecisionStatus.SUPERSEDED
            existing.updated_at = _now()
    conversation.decisions.append(decision)
    refresh_conversation_state(conversation)
    return decision


def _clean_payload(payload: dict[str, Any], model_cls: type) -> dict[str, Any]:
    """Strip unknown keys and ensure required fields have fallback values."""
    model_fields = set(model_cls.model_fields.keys())
    raw_desc = payload.pop("_raw", None)
    cleaned = {k: v for k, v in payload.items() if k in model_fields}
    for name, field_info in model_cls.model_fields.items():
        if name not in cleaned and field_info.is_required():
            cleaned[name] = raw_desc or "" if field_info.annotation is str or field_info.annotation == str else ""
    return cleaned


def apply_decision_proposal(conversation: Conversation, decision: DecisionRecord) -> DecisionRecord:
    contract = conversation.contract
    payload = decision.after or {}

    if decision.type == DecisionType.DATASET:
        current = contract.dataset.model_dump(mode="json") if contract.dataset else {}
        decision.before = decision.before or current
        contract.dataset = DatasetAsset(**_clean_payload(payload, DatasetAsset))
        contract.dataset.lock_state = LockState.PROPOSED
    elif decision.type == DecisionType.PREPROCESSING:
        current = contract.preprocessing.model_dump(mode="json") if contract.preprocessing else {}
        decision.before = decision.before or current
        contract.preprocessing = PreprocessingPlan(**_clean_payload(payload, PreprocessingPlan))
        contract.preprocessing.lock_state = LockState.PROPOSED
    elif decision.type == DecisionType.EVALUATION:
        current = contract.eval_spec.model_dump(mode="json") if contract.eval_spec else {}
        decision.before = decision.before or current
        contract.eval_spec = EvalSpec(**_clean_payload(payload, EvalSpec))
        contract.eval_spec.lock_state = LockState.PROPOSED
    elif decision.type == DecisionType.MODEL:
        current = contract.model_spec.model_dump(mode="json") if contract.model_spec else {}
        decision.before = decision.before or current
        contract.model_spec = ModelSpec(**_clean_payload(payload, ModelSpec))
        contract.model_spec.lock_state = LockState.PROPOSED
    elif decision.type == DecisionType.TARGET:
        decision.before = decision.before or _decision_snapshot(contract, decision.type)
        for key in ("target_threshold", "threshold_direction", "budget_minutes", "max_experiments"):
            if key in payload:
                setattr(contract, key, payload[key])
    _PRESERVE_AFTER_TYPES = {DecisionType.GENERAL, DecisionType.POST_RUN_RECOMMENDATION}
    if decision.type not in _PRESERVE_AFTER_TYPES:
        decision.after = _decision_snapshot(contract, decision.type)
    else:
        decision.after = payload
    return add_or_replace_decision(conversation, decision)


def apply_decision_action(
    conversation: Conversation,
    decision_id: str,
    action: str,
    note: str = "",
) -> DecisionRecord:
    match = next((d for d in conversation.decisions if d.id == decision_id), None)
    if not match:
        raise ValueError("Decision not found")

    match.updated_at = _now()
    if note:
        match.user_note = note.strip()

    if action == "approve":
        match.status = DecisionStatus.APPROVED
        _set_component_state(conversation.contract, match.type, LockState.FROZEN)
        if match.type == DecisionType.TARGET:
            conversation.contract.frozen = False
        sync_phase_states(conversation.contract)
        if conversation.state in {ConversationState.FAILED, ConversationState.STOPPED}:
            conversation.state = ConversationState.DRAFT
    elif action == "revoke":
        match.status = DecisionStatus.PROPOSED
        _set_component_state(conversation.contract, match.type, LockState.PROPOSED)
        _invalidate_for_decision(conversation, match.type, note or f"{match.title or match.type.value} revoked")
    elif action in {"reject", "request_changes"}:
        match.status = DecisionStatus.REJECTED if action == "reject" else DecisionStatus.NEEDS_INPUT
        _set_component_state(conversation.contract, match.type, LockState.PROPOSED)
        _invalidate_for_decision(conversation, match.type, note or f"{match.title or match.type.value} reopened")
    elif action == "provide_input":
        match.status = DecisionStatus.SUPERSEDED
    elif action == "complete":
        match.status = DecisionStatus.COMPLETED
    else:
        raise ValueError(f"Unsupported decision action: {action}")

    refresh_conversation_state(conversation)
    return match


def latest_post_run_decision(conversation: Conversation) -> DecisionRecord | None:
    for decision in reversed(conversation.decisions):
        if decision.type == DecisionType.POST_RUN_RECOMMENDATION:
            return decision
    return None


def _approved_components(conversation: Conversation) -> dict[str, bool]:
    contract = conversation.contract
    return {
        "dataset": bool(contract.dataset and contract.dataset.lock_state == LockState.FROZEN),
        "preprocessing": bool(
            contract.preprocessing is not None
            and contract.preprocessing.lock_state == LockState.FROZEN
        ),
        "model": bool(contract.model_spec and contract.model_spec.lock_state == LockState.FROZEN),
        "evaluation": bool(contract.eval_spec and contract.eval_spec.lock_state == LockState.FROZEN),
        "target": contract.target_threshold is not None,
    }


def compute_readiness(conversation: Conversation, *, is_processing: bool = False) -> Readiness:
    sync_phase_states(conversation.contract)
    contract = conversation.contract
    decisions = conversation.decisions
    blockers: list[str] = []
    next_actions: list[NextAction] = []
    approved = _approved_components(conversation)

    open_decisions = [
        d for d in decisions
        if d.status in {DecisionStatus.PROPOSED, DecisionStatus.NEEDS_INPUT, DecisionStatus.BLOCKED}
    ]
    active_decision = open_decisions[-1] if open_decisions else None

    if not approved["dataset"]:
        blockers.append("Approve a dataset decision before execution can start.")
    if not approved["model"]:
        blockers.append("Approve a model decision before training can run.")
    if not approved["evaluation"]:
        blockers.append("Approve an evaluation decision before training can run.")
    if not approved["target"]:
        blockers.append("Define a target threshold or stop condition.")

    if active_decision:
        next_actions.append(NextAction(
            action_type=NextActionType.APPROVE_DECISION,
            label=f"Approve {active_decision.title or active_decision.type.value}",
            description="Accept the current proposed change and update the spec.",
            decision_id=active_decision.id,
            enabled=active_decision.status == DecisionStatus.PROPOSED,
        ))
        next_actions.append(NextAction(
            action_type=NextActionType.REVISE_DECISION,
            label=f"Revise {active_decision.title or active_decision.type.value}",
            description="Request changes or provide the missing detail needed to move forward.",
            decision_id=active_decision.id,
        ))

    preprocessing_phase = contract.phases[PhaseName.PREPROCESSING.value]
    featurization_phase = contract.phases[PhaseName.FEATURIZATION.value]
    training_phase = contract.phases[PhaseName.TRAINING.value]
    evaluation_phase = contract.phases[PhaseName.EVALUATION.value]

    preprocessing_ready = approved["dataset"] and approved["preprocessing"]
    upstream_ready = (
        preprocessing_ready
        and approved["model"]
        and approved["evaluation"]
        and approved["target"]
    )

    preprocessing_invalid = preprocessing_phase.status in {
        PhaseExecutionStatus.PENDING,
        PhaseExecutionStatus.INVALIDATED,
        PhaseExecutionStatus.FAILED,
    }
    features_invalid = featurization_phase.status in {
        PhaseExecutionStatus.PENDING,
        PhaseExecutionStatus.INVALIDATED,
        PhaseExecutionStatus.FAILED,
    }
    plan_complete = bool(
        contract.preprocessing
        and contract.preprocessing.lock_state == LockState.FROZEN
    )
    no_pending = not open_decisions
    all_components_proposed = (
        contract.model_spec is not None
        and contract.eval_spec is not None
        and contract.target_threshold is not None
    )
    can_run_preprocessing = (
        no_pending
        and approved["dataset"]
        and plan_complete
        and all_components_proposed
        and (preprocessing_invalid or features_invalid)
    )
    can_run_training = no_pending and upstream_ready and not can_run_preprocessing

    if can_run_preprocessing:
        next_actions.append(NextAction(
            action_type=NextActionType.RUN_PREPROCESSING,
            label="Run preprocessing pipeline",
            description="Materialize reusable preprocessing and feature artifacts.",
            phase_name=PhaseName.PREPROCESSING.value,
            enabled=not open_decisions,
        ))
    if can_run_training:
        next_actions.append(NextAction(
            action_type=NextActionType.RUN_TRAINING,
            label="Run training and evaluation",
            description="Launch the execution phases using the approved spec.",
            phase_name=PhaseName.TRAINING.value,
            enabled=not open_decisions,
        ))

    latest_post_run = latest_post_run_decision(conversation)
    if (
        conversation.state != ConversationState.COMPLETED
        and latest_post_run
        and latest_post_run.status in {
        DecisionStatus.PROPOSED,
        DecisionStatus.APPROVED,
        DecisionStatus.COMPLETED,
        }
    ):
        next_actions.append(NextAction(
            action_type=NextActionType.REVIEW_RESULTS,
            label="Review latest run recommendation",
            description="Inspect the latest outcome and decide whether to iterate or finalize.",
            decision_id=latest_post_run.id,
        ))

    can_finalize = (
        evaluation_phase.status == PhaseExecutionStatus.COMPLETED
        and latest_post_run is not None
        and conversation.state != ConversationState.COMPLETED
    )
    if can_finalize:
        next_actions.append(NextAction(
            action_type=NextActionType.FINALIZE,
            label="Finalize current result",
            description="Stop iterating and treat the latest verified result as the current best outcome.",
        ))

    if not next_actions:
        next_actions.append(NextAction(
            action_type=NextActionType.STOP,
            label="Stop work",
            description="Explicitly end the current session.",
        ))

    if conversation.state == ConversationState.RUNNING:
        current_stage = ConversationState.RUNNING
    elif conversation.state == ConversationState.PAUSED:
        current_stage = ConversationState.PAUSED
    elif active_decision and active_decision.status == DecisionStatus.NEEDS_INPUT:
        current_stage = ConversationState.NEEDS_INPUT
    elif latest_post_run and evaluation_phase.status == PhaseExecutionStatus.COMPLETED:
        current_stage = ConversationState.REVIEWING_RESULTS
    elif can_run_training:
        current_stage = ConversationState.READY_TO_TRAIN
    elif can_run_preprocessing:
        current_stage = ConversationState.READY_TO_PREPROCESS
    else:
        current_stage = ConversationState.DRAFT

    reusable_artifacts = []
    for phase_name in (PhaseName.PREPROCESSING.value, PhaseName.FEATURIZATION.value):
        phase = contract.phases[phase_name]
        if phase.status == PhaseExecutionStatus.COMPLETED and phase.artifact_refs:
            reusable_artifacts.append(phase.label or phase_name)

    proposed_count = sum(1 for d in open_decisions if d.status == DecisionStatus.PROPOSED)
    needs_input_count = sum(1 for d in open_decisions if d.status == DecisionStatus.NEEDS_INPUT)
    pending_review_count = proposed_count + needs_input_count

    if current_stage == ConversationState.NEEDS_INPUT:
        summary = "Waiting for missing user input."
    elif current_stage == ConversationState.READY_TO_PREPROCESS:
        summary = "The approved spec is ready for preprocessing."
    elif current_stage == ConversationState.READY_TO_TRAIN:
        summary = "The approved spec is ready for training and evaluation."
    elif current_stage == ConversationState.REVIEWING_RESULTS:
        summary = "A run finished and the next recommendation is ready for review."
    elif current_stage == ConversationState.RUNNING:
        summary = "Execution is in progress."
    elif pending_review_count > 0:
        label = "decision" if pending_review_count == 1 else "decisions"
        summary = f"{pending_review_count} {label} pending your review."
    elif blockers:
        summary = blockers[0]
    else:
        summary = "Keep refining the spec."

    if conversation.state in {ConversationState.FAILED, ConversationState.STOPPED}:
        agent_status, agent_detail = "error", "Something went wrong."
    elif conversation.state == ConversationState.RUNNING:
        agent_status, agent_detail = "running", "Execution in progress."
    elif conversation.state == ConversationState.PAUSED:
        agent_status, agent_detail = "paused", "Execution paused."
    elif is_processing:
        agent_status, agent_detail = "thinking", "Agent is working..."
    elif needs_input_count > 0:
        q = "question" if needs_input_count == 1 else "questions"
        agent_status, agent_detail = "awaiting_input", f"{needs_input_count} {q} need your answer."
    elif proposed_count > 0:
        d = "decision" if proposed_count == 1 else "decisions"
        agent_status, agent_detail = "awaiting_approval", f"{proposed_count} {d} pending your review."
    elif conversation.state == ConversationState.REVIEWING_RESULTS:
        agent_status, agent_detail = "reviewing", "Results available for review."
    elif conversation.state == ConversationState.COMPLETED:
        agent_status, agent_detail = "completed", "Pipeline complete."
    elif can_run_preprocessing or can_run_training:
        label = "preprocessing" if can_run_preprocessing else "training"
        agent_status, agent_detail = "ready_to_run", f"Ready to start {label}."
    elif not conversation.messages:
        agent_status, agent_detail = "idle", "Start by describing your ML task."
    else:
        agent_status, agent_detail = "idle", summary

    return Readiness(
        current_stage=current_stage,
        can_run_preprocessing=can_run_preprocessing,
        can_run_training=can_run_training,
        can_finalize=can_finalize,
        blocking_reasons=blockers,
        next_actions=next_actions,
        reusable_artifacts=reusable_artifacts,
        active_decision_id=active_decision.id if active_decision else None,
        pending_decision_count=pending_review_count,
        summary=summary,
        agent_status=agent_status,
        agent_status_detail=agent_detail,
    )


def refresh_conversation_state(conversation: Conversation, *, is_processing: bool = False) -> Conversation:
    conversation.readiness = compute_readiness(conversation, is_processing=is_processing)
    _STICKY_STATES = {
        ConversationState.RUNNING,
        ConversationState.PAUSED,
        ConversationState.FAILED,
        ConversationState.STOPPED,
        ConversationState.COMPLETED,
    }
    if conversation.state not in _STICKY_STATES:
        conversation.state = conversation.readiness.current_stage
    return conversation
