"""Tests for Phase 2 self-learning loop orchestration."""

from __future__ import annotations

import pytest

from logos_adk.agent import Agent
from logos_adk.learning import (
    FeedbackRecord,
    PromptRegistry,
    PromptCandidateGenerator,
    ReflectionMemory,
    SelfLearningLoop,
    SelfLearningLoopPolicy,
    SQLiteLearningStore,
)
from logos_adk.learning.improvement.models import ReflectionLesson
from logos_adk.learning.models import RunRecord
from logos_adk.providers.mock import MockProvider
from logos_adk.types import Response, Usage


class CandidateAwareProvider(MockProvider):
    async def generate(self, messages, tools=None, output_schema=None, **kwargs):
        system_prompt = next(
            (message.content for message in messages if message.role == "system"), ""
        )
        content = "candidate-good" if "candidate-good" in system_prompt else "candidate-bad"
        return Response(
            content=content,
            tool_calls=[],
            usage=Usage(input_tokens=20, output_tokens=8),
            raw={},
        )


def _seed_bad_case(
    store: SQLiteLearningStore, *, trace_id: str, corrected_output: str = "candidate-good"
) -> None:
    store.create_run(
        RunRecord(
            id=f"run-{trace_id}",
            trace_id=trace_id,
            request_id=f"req-{trace_id}",
            agent_name="bot",
            session_id="s-1",
            workspace_id="acme",
            prompt_version="v1",
            model="mock:model",
            config_fingerprint="cfg-v1",
            input_text="write reply",
            output_text="candidate-bad",
            metadata={"prompt_name": "bot"},
        )
    )
    store.create_feedback(
        FeedbackRecord(
            id=f"fb-{trace_id}",
            kind="edited_answer",
            trace_id=trace_id,
            run_id=f"run-{trace_id}",
            request_id=f"req-{trace_id}",
            session_id="s-1",
            workspace_id="acme",
            agent_name="bot",
            prompt_version="v1",
            model="mock:model",
            config_fingerprint="cfg-v1",
            corrected_output=corrected_output,
            failure_types=["hallucination", "style_miss"],
        )
    )


@pytest.mark.asyncio
async def test_self_learning_loop_supports_dry_run(tmp_path):
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    registry = PromptRegistry(str(tmp_path / "prompts.db"))
    reflection_memory = ReflectionMemory(path=str(tmp_path / "reflection.db"))
    registry.ensure_prompt("bot")
    baseline = registry.create_version("bot", content="baseline prompt")
    registry.promote("bot", version=baseline.version)
    _seed_bad_case(store, trace_id="trace-dry-run")

    loop = SelfLearningLoop(store, registry, reflection_memory=reflection_memory)
    agent = Agent("bot", model=CandidateAwareProvider(), system="fallback")
    result = await loop.run_once(
        agent,
        prompt_name="bot",
        workspace_id="acme",
        policy=SelfLearningLoopPolicy(
            min_sample_size=1, min_improvement=0.01, allow_auto_promote=False
        ),
        dry_run=True,
    )

    assert result.action is not None
    assert result.action.kind == "experiment"
    assert result.action.status == "dry_run"
    assert result.created_version is None
    assert len(registry.list_versions("bot")) == 1
    assert "dry_run" in result.audit_trail


@pytest.mark.asyncio
async def test_self_learning_loop_creates_prompt_experiment(tmp_path):
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    registry = PromptRegistry(str(tmp_path / "prompts.db"))
    reflection_memory = ReflectionMemory(path=str(tmp_path / "reflection.db"))
    registry.ensure_prompt("bot")
    baseline = registry.create_version("bot", content="baseline prompt")
    registry.promote("bot", version=baseline.version)
    _seed_bad_case(store, trace_id="trace-experiment")

    loop = SelfLearningLoop(store, registry, reflection_memory=reflection_memory)
    agent = Agent("bot", model=CandidateAwareProvider(), system="fallback")
    result = await loop.run_once(
        agent,
        prompt_name="bot",
        workspace_id="acme",
        policy=SelfLearningLoopPolicy(
            min_sample_size=1, min_improvement=0.01, allow_auto_promote=False
        ),
    )

    assert result.action is not None
    assert result.action.kind == "experiment"
    assert result.action.status == "executed"
    assert result.created_version is not None
    assert result.experiment is not None
    experiments = registry.list_experiments("bot", status="active")
    assert len(experiments) == 1
    assert experiments[0].baseline_version == baseline.version
    assert experiments[0].challenger_versions == [result.created_version["version"]]


@pytest.mark.asyncio
async def test_self_learning_loop_can_auto_promote(tmp_path):
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    registry = PromptRegistry(str(tmp_path / "prompts.db"))
    reflection_memory = ReflectionMemory(path=str(tmp_path / "reflection.db"))
    registry.ensure_prompt("bot")
    baseline = registry.create_version("bot", content="baseline prompt")
    registry.promote("bot", version=baseline.version)
    _seed_bad_case(store, trace_id="trace-promote")

    loop = SelfLearningLoop(store, registry, reflection_memory=reflection_memory)
    agent = Agent("bot", model=CandidateAwareProvider(), system="fallback")
    result = await loop.run_once(
        agent,
        prompt_name="bot",
        workspace_id="acme",
        policy=SelfLearningLoopPolicy(
            min_sample_size=1, min_improvement=0.01, allow_auto_promote=True
        ),
    )

    assert result.action is not None
    assert result.action.kind == "promote"
    assert result.action.status == "executed"
    assert result.deployment is not None
    resolved = registry.resolve("bot", workspace_id="acme")
    assert resolved is not None
    assert resolved.version == result.created_version["version"]


def test_prompt_candidate_generator_persists_generation_report(tmp_path):
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    registry = PromptRegistry(str(tmp_path / "prompts.db"))
    reflection_memory = ReflectionMemory(path=str(tmp_path / "reflection.db"))
    registry.ensure_prompt("bot")
    baseline = registry.create_version("bot", content="baseline prompt")
    registry.promote("bot", version=baseline.version, workspace_id="acme")
    _seed_bad_case(store, trace_id="trace-report")
    reflection_memory.save_lesson(
        ReflectionLesson(
            id="lesson-1",
            workspace_id="acme",
            agent_name="bot",
            category="prompt_patch",
            lesson="Recurring hallucination failures",
            rationale="Multiple corrected outputs required stronger grounding.",
            suggested_action="Require grounded answers and explicit uncertainty.",
            confidence=0.8,
            source_run_ids=["run-trace-report"],
        )
    )

    generator = PromptCandidateGenerator(registry, store=store, reflection_memory=reflection_memory)
    report = generator.generate_report(prompt_name="bot", workspace_id="acme", max_candidates=5)

    assert report.based_on_version == baseline.version
    assert report.corrected_output_count >= 1
    assert len(report.candidates) >= 2
    recommendations = generator.persist_candidates(report)
    saved = registry.list_recommendations("bot")
    assert len(recommendations) == len(report.candidates)
    assert any(item.kind == "generated_candidate" for item in saved)
    assert any(item.metadata.get("strategy") == "style_correction" for item in saved)


@pytest.mark.asyncio
async def test_self_learning_loop_records_generation_report_and_candidate_provenance(tmp_path):
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    registry = PromptRegistry(str(tmp_path / "prompts.db"))
    reflection_memory = ReflectionMemory(path=str(tmp_path / "reflection.db"))
    registry.ensure_prompt("bot")
    baseline = registry.create_version("bot", content="baseline prompt")
    registry.promote("bot", version=baseline.version, workspace_id="acme")
    _seed_bad_case(store, trace_id="trace-phase3")
    reflection_memory.save_lesson(
        ReflectionLesson(
            id="lesson-2",
            workspace_id="acme",
            agent_name="bot",
            category="prompt_patch",
            lesson="Use corrected answers as exemplars",
            rationale="Human edits consistently improve style.",
            suggested_action="Prefer concise factual tone aligned with corrected outputs.",
            confidence=0.9,
            source_run_ids=["run-trace-phase3"],
        )
    )

    loop = SelfLearningLoop(store, registry, reflection_memory=reflection_memory)
    agent = Agent("bot", model=CandidateAwareProvider(), system="fallback")
    result = await loop.run_once(
        agent,
        prompt_name="bot",
        workspace_id="acme",
        policy=SelfLearningLoopPolicy(
            min_sample_size=1, min_improvement=0.01, allow_auto_promote=False
        ),
    )

    assert result.generation_report is not None
    assert result.generation_report["based_on_version"] == baseline.version
    assert len(result.generation_report["candidates"]) >= 1
    assert any(
        item["strategy"] in {"patch", "style_correction", "rewrite"}
        for item in result.generation_report["candidates"]
    )
    created = registry.get_version("bot", result.created_version["version"])
    assert created is not None
    assert created.notes is not None
    assert "generated_candidates" in created.notes
    recommendations = registry.list_recommendations("bot")
    assert any(item.kind == "generated_candidate" for item in recommendations)


@pytest.mark.asyncio
async def test_self_learning_loop_supports_learned_reward_scorer_mode(tmp_path):
    store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    registry = PromptRegistry(str(tmp_path / "prompts.db"))
    reflection_memory = ReflectionMemory(path=str(tmp_path / "reflection.db"))
    registry.ensure_prompt("bot")
    baseline = registry.create_version("bot", content="baseline prompt")
    registry.promote("bot", version=baseline.version, workspace_id="acme")
    _seed_bad_case(store, trace_id="trace-learned-1")
    _seed_bad_case(store, trace_id="trace-learned-2")

    loop = SelfLearningLoop(store, registry, reflection_memory=reflection_memory)
    agent = Agent("bot", model=CandidateAwareProvider(), system="fallback")
    result = await loop.run_once(
        agent,
        prompt_name="bot",
        workspace_id="acme",
        policy=SelfLearningLoopPolicy(
            min_sample_size=1,
            min_improvement=0.01,
            allow_auto_promote=False,
            reward_scorer_mode="learned",
            reward_training_limit=50,
        ),
    )

    assert result.reward_scorer_mode == "learned"
    assert result.historical_reward_score >= 0.0
    assert "reward_scorer_mode=learned" in result.audit_trail
    assert result.created_version is not None
    created = registry.get_version("bot", result.created_version["version"])
    assert created is not None
    assert created.notes is not None
    assert '"reward_scorer_mode": "learned"' in created.notes
