from __future__ import annotations

from pathlib import Path

import pytest

from epochler.agents import planner
from epochler.contract.engine import mark_phase_completed
from epochler.contract.models import (
    ChatMessage,
    Conversation,
    DatasetAsset,
    DecisionRecord,
    DecisionType,
    EvalSpec,
    HardwareProfile,
    LockState,
    ModelSpec,
    PhaseName,
    PreprocessingPlan,
)
from epochler.decisioning import apply_decision_action, apply_decision_proposal, refresh_conversation_state
from epochler.runner.runner import ExperimentRunner
from epochler.storage import store


def _decision(decision_type: DecisionType, title: str, after: dict) -> DecisionRecord:
    return DecisionRecord(
        type=decision_type,
        title=title,
        summary=title,
        rationale="test",
        after=after,
    )


def test_readiness_flows_from_decisions_to_phases() -> None:
    conv = Conversation()

    apply_decision_proposal(conv, _decision(DecisionType.DATASET, "dataset", {
        "path": "/tmp/data.csv",
        "format": "csv",
        "schema_info": {},
        "target_column": "label",
    }))
    apply_decision_action(conv, conv.decisions[-1].id, "approve")

    apply_decision_proposal(conv, _decision(DecisionType.PREPROCESSING, "pre", {
        "source": "test",
        "transforms": [{"name": "scale"}],
        "reproducibility_label": "exact",
    }))
    apply_decision_action(conv, conv.decisions[-1].id, "approve")

    apply_decision_proposal(conv, _decision(DecisionType.EVALUATION, "eval", {
        "primary_metric": "accuracy",
        "secondary_metrics": ["f1"],
        "split_strategy": "stratified",
        "num_seeds": 1,
        "constraints": {},
    }))
    apply_decision_action(conv, conv.decisions[-1].id, "approve")

    apply_decision_proposal(conv, _decision(DecisionType.MODEL, "model", {
        "allowed_families": ["random_forest"],
        "search_space": {"max_depth": [4, 8]},
        "hardware_constraints": {},
    }))
    apply_decision_action(conv, conv.decisions[-1].id, "approve")

    apply_decision_proposal(conv, _decision(DecisionType.TARGET, "target", {
        "target_threshold": 0.8,
        "threshold_direction": "higher_is_better",
        "budget_minutes": 5,
        "max_experiments": 2,
    }))
    apply_decision_action(conv, conv.decisions[-1].id, "approve")

    refresh_conversation_state(conv)
    assert conv.readiness.can_run_preprocessing is True
    assert conv.readiness.can_run_training is False

    mark_phase_completed(conv.contract, PhaseName.PREPROCESSING.value, [])
    mark_phase_completed(conv.contract, PhaseName.FEATURIZATION.value, [])
    refresh_conversation_state(conv)
    assert conv.readiness.can_run_training is True

    post_run = DecisionRecord(
        type=DecisionType.POST_RUN_RECOMMENDATION,
        title="Post-run recommendation",
        summary="Review the latest run.",
        rationale="test",
        after={"recommended_next_action": "finalize"},
    )
    apply_decision_proposal(conv, post_run)
    mark_phase_completed(conv.contract, PhaseName.EVALUATION.value, [])
    refresh_conversation_state(conv)

    assert conv.readiness.can_finalize is True
    assert any(action.action_type.value == "finalize" for action in conv.readiness.next_actions)


@pytest.mark.asyncio
async def test_planner_falls_back_when_provider_fails(monkeypatch: pytest.MonkeyPatch) -> None:
    conv = Conversation()
    conv.messages.append(ChatMessage(role="user", content="/tmp/data.csv"))

    class FakeMessages:
        async def create(self, **kwargs):
            raise RuntimeError("provider credits exhausted")

    class FakeClient:
        def __init__(self, **kwargs):
            self.messages = FakeMessages()

    monkeypatch.setattr(planner, "get_settings", lambda: type("Settings", (), {
        "anthropic_api_key": "present",
        "anthropic_model": "dummy",
    })())
    monkeypatch.setattr(planner.anthropic, "AsyncAnthropic", FakeClient)

    response = await planner.respond(conv, HardwareProfile(cpu_count=2, ram_total_gb=8.0, gpu_available=False))

    assert response["decisions"]
    assert response["decisions"][0]["type"] == DecisionType.DATASET.value
    assert "dataset" in response["assistant_message"].lower()


@pytest.mark.asyncio
async def test_runner_persists_versioned_artifacts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(store, "_data_dir", lambda: tmp_path)
    monkeypatch.setattr(store, "workspaces_dir", lambda: tmp_path / "workspaces")

    dataset_path = tmp_path / "dataset.csv"
    dataset_path.write_text("feature,label\n1,0\n2,1\n", encoding="utf-8")

    preprocessing_code = (
        "import json\n"
        "import os\n"
        "from pathlib import Path\n"
        "out_dir = Path(os.environ['EPOCHLER_OUTPUT_DIR'])\n"
        "out_dir.mkdir(parents=True, exist_ok=True)\n"
        "print(json.dumps({'artifacts': [{'name': 'preprocessed_dataset', 'path': str(out_dir), 'artifact_type': 'preprocessed_dataset'}]}))\n"
    )
    training_code = (
        "import json\n"
        "import os\n"
        "from pathlib import Path\n"
        "out_dir = Path(os.environ['EPOCHLER_OUTPUT_DIR'])\n"
        "out_dir.mkdir(parents=True, exist_ok=True)\n"
        "model_path = out_dir / 'model.txt'\n"
        "model_path.write_text('placeholder-model', encoding='utf-8')\n"
        "print(json.dumps({'metrics': {'accuracy': 0.5}, 'artifacts': [{'name': 'model_checkpoint', 'path': str(model_path), 'artifact_type': 'model_checkpoint'}]}))\n"
    )

    async def fake_generate_phase_script(*args, **kwargs):
        phase_name = kwargs["phase_name"]
        code = preprocessing_code if phase_name == PhaseName.PREPROCESSING.value else training_code
        return {"code": code, "description": "generated"}

    async def fake_generate_iteration(contract, current_code, journal, hardware, input_artifacts):
        return {
            "code": training_code,
            "hypothesis": "baseline",
            "description": "baseline",
            "search_queries": [],
        }

    from epochler.agents import coder
    import epochler.runner.runner as runner_module

    monkeypatch.setattr(coder, "generate_phase_script", fake_generate_phase_script)
    monkeypatch.setattr(coder, "generate_iteration", fake_generate_iteration)
    monkeypatch.setattr(
        runner_module,
        "detect_hardware",
        lambda: HardwareProfile(cpu_count=2, ram_total_gb=2.0, gpu_available=False),
    )

    contract = Conversation().contract
    contract.dataset = DatasetAsset(path=str(dataset_path), format="csv", target_column="label", lock_state=LockState.FROZEN)
    contract.preprocessing = PreprocessingPlan(source="test", transforms=[], reproducibility_label="exact", lock_state=LockState.FROZEN)
    contract.eval_spec = EvalSpec(primary_metric="accuracy", secondary_metrics=["f1"], split_strategy="stratified", num_seeds=1, constraints={}, lock_state=LockState.FROZEN)
    contract.model_spec = ModelSpec(allowed_families=["random_forest"], search_space={}, hardware_constraints={}, lock_state=LockState.FROZEN)
    contract.target_threshold = 0.5
    contract.threshold_direction = "higher_is_better"
    contract.max_experiments = 1
    contract.budget_minutes = 1

    events: list[dict] = []

    async def broadcast(event: dict) -> None:
        events.append(event)

    async def check_paused() -> bool:
        return False

    runner = ExperimentRunner(contract, broadcast, check_paused)
    journal = await runner.run()

    assert journal.experiments
    assert journal.experiments[0].metrics["accuracy"] == 0.5
    versions = store.list_artifact_versions(contract.id, PhaseName.TRAINING.value)
    assert versions
