from __future__ import annotations

import asyncio
import json
from pathlib import Path
from types import SimpleNamespace

from beaver.engine.providers.base import LLMProvider, LLMResponse
from beaver.engine.providers.factory import ProviderBundle
from beaver.engine.session import SessionManager
from beaver.memory.runs import RunMemoryStore, RunRecord
from beaver.memory.skills import SkillLearningCandidate, SkillLearningStore
from beaver.skills.authoring.format import is_canonical_skill_body
from beaver.skills.drafts import DraftService
from beaver.skills.learning import (
    EvidenceSelector,
    SkillDraftSynthesizer,
    SkillLearningPipelineService,
    SkillLearningService,
    SkillLearningWorker,
    SkillLearningWorkerConfig,
)
from beaver.skills.publisher import SkillPublisher
from beaver.skills.reviews import ReviewService
from beaver.skills.specs import SkillSpecStore


class JsonProvider(LLMProvider):
    def __init__(self, payload: dict | None = None, *, fail: bool = False) -> None:
        super().__init__()
        self.payload = payload or {
            "frontmatter": {"description": "Generated skill", "tools": []},
            "content": "# Generated\n\nUse the learned workflow.",
            "change_reason": "learned",
        }
        self.fail = fail

    async def chat(self, messages: list[dict], tools: list[dict] | None = None, model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7) -> LLMResponse:
        if self.fail:
            raise RuntimeError("provider failed")
        return LLMResponse(content=json.dumps(self.payload), model=model)

    def get_default_model(self) -> str:
        return "stub"


def _bundle(provider: LLMProvider) -> ProviderBundle:
    runtime = SimpleNamespace(model="stub", provider_name="stub")
    return ProviderBundle(main_runtime=runtime, main_provider=provider)  # type: ignore[arg-type]


class FakeReplayRunner:
    def __init__(self) -> None:
        self.requests = []

    async def run_arm(self, request):
        self.requests.append(request)
        return {
            "case_id": request.case_id,
            "arm": request.arm,
            "session_id": "session-replay",
            "run_id": f"{request.arm}-run",
            "task_text": request.task_text,
            "finish_reason": "stop",
            "final_answer": "debug deployment startup done",
            "tool_calls": [
                {
                    "tool_name": "echo",
                    "mode": "executed",
                    "arguments": {"text": "ok"},
                    "result": {"success": True, "content": "ok"},
                }
            ],
            "artifacts": [],
            "side_effects": [],
        }


def _pipeline(tmp_path: Path) -> SkillLearningPipelineService:
    spec_store = SkillSpecStore(tmp_path)
    run_store = RunMemoryStore(tmp_path / "memory" / "runs")
    learning_store = SkillLearningStore(tmp_path / "memory" / "skills")
    run_store.append_run_record(
        RunRecord(
            run_id="run-1",
            session_id="session-1",
            task_text="debug deployment startup",
            started_at="start",
            ended_at="end",
            success=True,
            finish_reason="stop",
        )
    )
    learning_store.record_learning_candidate(
        SkillLearningCandidate(
            candidate_id="candidate-1",
            kind="new_skill",
            source_run_ids=["run-1"],
            source_session_ids=["session-1"],
            related_skill_names=[],
            reason="repeat success",
            priority=10,
            confidence=0.9,
        )
    )
    draft_service = DraftService(spec_store)
    learning_service = SkillLearningService(
        run_store=run_store,
        learning_store=learning_store,
        draft_service=draft_service,
        evidence_selector=EvidenceSelector(run_store),
        synthesizer=SkillDraftSynthesizer(),
    )
    return SkillLearningPipelineService(
        learning_store=learning_store,
        learning_service=learning_service,
        draft_service=draft_service,
        review_service=ReviewService(spec_store),
        publisher=SkillPublisher(spec_store),
    )


def test_worker_synthesizes_open_candidate_without_publish(tmp_path: Path) -> None:
    pipeline = _pipeline(tmp_path)
    worker = SkillLearningWorker(
        pipeline=pipeline,
        provider_bundle_factory=lambda: _bundle(JsonProvider()),
        config=SkillLearningWorkerConfig(max_drafts_per_run=5, max_retries=3, interval_seconds=1),
    )

    result = asyncio.run(worker.run_once())
    candidate = pipeline.get_candidate("candidate-1")

    assert result.succeeded == 1
    assert candidate.status == "draft_ready"
    assert candidate.draft_id
    assert pipeline.list_drafts(candidate.draft_skill_name)[0].status == "draft"


def test_worker_evaluates_draft_with_replay_runner_when_available(tmp_path: Path) -> None:
    pipeline = _pipeline(tmp_path)
    replay_runner = FakeReplayRunner()
    worker = SkillLearningWorker(
        pipeline=pipeline,
        provider_bundle_factory=lambda: _bundle(JsonProvider()),
        replay_runner_factory=lambda: replay_runner,
        config=SkillLearningWorkerConfig(max_drafts_per_run=5, max_retries=3, interval_seconds=1),
    )

    result = asyncio.run(worker.run_once())
    candidate = pipeline.get_candidate("candidate-1")
    draft = pipeline.get_draft(candidate.draft_skill_name or "", candidate.draft_id or "")
    report = pipeline.get_eval_report(draft.skill_name, draft.draft_id)

    assert result.succeeded == 1
    assert report is not None
    assert report.mode == "replay"
    assert report.case_reports
    assert replay_runner.requests


def test_worker_retries_and_marks_failed_after_limit(tmp_path: Path) -> None:
    pipeline = _pipeline(tmp_path)
    worker = SkillLearningWorker(
        pipeline=pipeline,
        provider_bundle_factory=lambda: _bundle(JsonProvider(fail=True)),
        config=SkillLearningWorkerConfig(max_drafts_per_run=5, max_retries=1, interval_seconds=1),
    )

    result = asyncio.run(worker.run_once())
    candidate = pipeline.get_candidate("candidate-1")

    assert result.failed == 1
    assert candidate.status == "failed"
    assert candidate.retry_count == 1
    assert "provider failed" in (candidate.last_error or "")


def test_synthesizer_fills_missing_tools_from_evidence(tmp_path: Path) -> None:
    pipeline = _pipeline(tmp_path)
    candidate = pipeline.get_candidate("candidate-1")
    provider = JsonProvider(
        payload={
            "frontmatter": {"description": "Generated skill"},
            "content": "# Generated\n\nUse the observed workflow.",
            "change_reason": "learned",
        }
    )
    packet = EvidenceSelector(pipeline.learning_service.run_store).build_evidence_packet(
        candidate.source_run_ids,
        candidate.source_session_ids,
    )
    packet.metadata["tool_names"] = ["web_fetch", "memory"]

    payload = asyncio.run(
        SkillDraftSynthesizer().synthesize_new_skill(candidate, packet, provider, "stub")
    )

    assert payload["frontmatter"]["tools"] == ["web_fetch", "memory"]
    assert is_canonical_skill_body(payload["content"])


def test_evidence_selector_records_run_tool_names(tmp_path: Path) -> None:
    run_store = RunMemoryStore(tmp_path / "memory" / "runs")
    run_store.append_run_record(
        RunRecord(
            run_id="run-1",
            session_id="session-1",
            task_text="research latest docs",
            started_at="start",
            ended_at="end",
            success=True,
            finish_reason="stop",
        )
    )
    session_manager = SessionManager(tmp_path)
    session_manager.ensure_session("session-1")
    session_manager.append_message(
        "session-1",
        run_id="run-1",
        role="system",
        event_type="tool_selection_snapshotted",
        event_payload={"tool_names": ["memory", "web_fetch"]},
        context_visible=False,
    )
    session_manager.append_message(
        "session-1",
        run_id="run-1",
        role="assistant",
        tool_calls=[{"id": "call-1", "function": {"name": "web_search"}}],
    )
    session_manager.append_message(
        "session-1",
        run_id="run-1",
        role="tool",
        tool_name="web_fetch",
        content="ok",
    )

    try:
        packet = EvidenceSelector(run_store, session_manager).build_evidence_packet(
            ["run-1"],
            ["session-1"],
        )
    finally:
        session_manager.close()

    assert packet.metadata["tool_names"] == ["web_search", "web_fetch"]
    assert packet.metadata["selected_tool_names"] == ["memory", "web_fetch"]


def test_worker_supersedes_candidate_when_active_draft_exists(tmp_path: Path) -> None:
    pipeline = _pipeline(tmp_path)
    pipeline.learning_store.record_learning_candidate(
        SkillLearningCandidate(
            candidate_id="candidate-2",
            kind="revise_skill",
            source_run_ids=["run-1"],
            source_session_ids=["session-1"],
            related_skill_names=["shared-skill"],
            reason="duplicate",
            status="draft_ready",
            draft_skill_name="shared-skill",
            draft_id="draft-existing",
        )
    )
    pipeline.learning_store.update_learning_candidate("candidate-1", related_skill_names=["shared-skill"])
    worker = SkillLearningWorker(
        pipeline=pipeline,
        provider_bundle_factory=lambda: _bundle(JsonProvider()),
        config=SkillLearningWorkerConfig(max_drafts_per_run=5, max_retries=3, interval_seconds=1),
    )

    result = asyncio.run(worker.run_once())

    assert result.skipped == 1
    assert pipeline.get_candidate("candidate-1").status == "superseded"
