"""Tests for P3 true improvement mechanisms."""

from __future__ import annotations

import json

import pytest

from logos_adk.learning import (
    CriticReviserLoop,
    HeuristicRewardScorer,
    LearnedPreferenceScorer,
    ModelRoutingLearningEngine,
    PatternMiner,
    PolicyLearner,
    ReflectionMemory,
    RetrievalLearningEngine,
    RootCauseAnalyzer,
    RewardModel,
    SQLiteLearningStore,
    ToolWorkflowLearningEngine,
    TrainingExportBuilder,
)
from logos_adk.learning.models import FeedbackRecord, LearningExample, OutcomeRecord, RunRecord
from logos_adk.rag import Chunk


def _run(
    *,
    run_id: str,
    trace_id: str,
    input_text: str,
    output_text: str,
    workspace_id: str = "acme",
    agent_name: str = "bot",
    prompt_version: str = "v1",
    model: str = "gpt-4o-mini",
    metadata: dict | None = None,
    usage: dict | None = None,
) -> RunRecord:
    return RunRecord(
        id=run_id,
        trace_id=trace_id,
        request_id=f"req-{run_id}",
        agent_name=agent_name,
        input_text=input_text,
        output_text=output_text,
        workspace_id=workspace_id,
        session_id="session-1",
        prompt_version=prompt_version,
        model=model,
        config_fingerprint="cfg-v1",
        usage=usage or {"input_tokens": 800, "output_tokens": 200},
        metadata=dict(metadata or {}),
    )


def _feedback(
    *,
    feedback_id: str,
    trace_id: str,
    run_id: str,
    kind: str,
    corrected_output: str | None = None,
    expected_output: str | None = None,
    rating: float | None = None,
    failure_types: list[str] | None = None,
) -> FeedbackRecord:
    return FeedbackRecord(
        id=feedback_id,
        kind=kind,  # type: ignore[arg-type]
        run_id=run_id,
        trace_id=trace_id,
        workspace_id="acme",
        agent_name="bot",
        prompt_version="v1",
        model="gpt-4o-mini",
        corrected_output=corrected_output,
        expected_output=expected_output,
        rating=rating,
        failure_types=list(failure_types or []),  # type: ignore[arg-type]
    )


def _outcome(
    *,
    outcome_id: str,
    trace_id: str,
    run_id: str,
    kind: str,
    value,
) -> OutcomeRecord:
    return OutcomeRecord(
        id=outcome_id,
        kind=kind,  # type: ignore[arg-type]
        run_id=run_id,
        trace_id=trace_id,
        workspace_id="acme",
        agent_name="bot",
        prompt_version="v1",
        model="gpt-4o-mini",
        value=value,
    )


def test_reward_model_aggregates_feedback_outcomes_and_penalties():
    reward_model = RewardModel()

    strong = LearningExample(
        run_id="run-good",
        trace_id="trace-good",
        agent_name="bot",
        input_text="Answer the user",
        output_text="Useful grounded answer",
        created_at="2026-03-18T00:00:00+00:00",
        workspace_id="acme",
        prompt_version="v2",
        model="gpt-4o-mini",
        feedback=[
            _feedback(
                feedback_id="fb-good", trace_id="trace-good", run_id="run-good", kind="thumbs_up"
            )
        ],
        outcomes=[
            _outcome(
                outcome_id="out-good",
                trace_id="trace-good",
                run_id="run-good",
                kind="conversion",
                value=True,
            )
        ],
        metadata={
            "usage": {"input_tokens": 500, "output_tokens": 100},
            "latency_ms": 300,
        },
    )
    weak = LearningExample(
        run_id="run-bad",
        trace_id="trace-bad",
        agent_name="bot",
        input_text="Answer the user",
        output_text="Unsafe answer",
        created_at="2026-03-18T00:00:00+00:00",
        workspace_id="acme",
        prompt_version="v2",
        model="claude-sonnet-4-20250514",
        feedback=[
            _feedback(
                feedback_id="fb-bad",
                trace_id="trace-bad",
                run_id="run-bad",
                kind="thumbs_down",
                failure_types=["policy_violation"],
            )
        ],
        outcomes=[
            _outcome(
                outcome_id="out-bad",
                trace_id="trace-bad",
                run_id="run-bad",
                kind="escalation",
                value=True,
            )
        ],
        failure_types=["policy_violation"],
        metadata={
            "usage": {"input_tokens": 50000, "output_tokens": 20000},
            "latency_ms": 9200,
        },
    )

    strong_score = reward_model.score_example(strong)
    weak_score = reward_model.score_example(weak)

    assert strong_score.score > weak_score.score
    assert {component.name for component in strong_score.components} == {
        "human_feedback",
        "business_outcome",
        "cost_penalty",
        "latency_penalty",
        "safety_penalty",
    }
    assert weak_score.score < 0.25


def test_reward_model_remains_backward_compatibility_alias():
    reward_model = RewardModel()
    heuristic = HeuristicRewardScorer()
    example = LearningExample(
        run_id="run-alias",
        trace_id="trace-alias",
        agent_name="bot",
        input_text="Respond safely",
        output_text="Grounded answer",
        created_at="2026-03-18T00:00:00+00:00",
        workspace_id="acme",
        prompt_version="v1",
        model="gpt-4o-mini",
        feedback=[
            _feedback(
                feedback_id="fb-alias", trace_id="trace-alias", run_id="run-alias", kind="thumbs_up"
            )
        ],
        outcomes=[
            _outcome(
                outcome_id="out-alias",
                trace_id="trace-alias",
                run_id="run-alias",
                kind="resolution",
                value=True,
            )
        ],
        metadata={"latency_ms": 120, "usage": {"input_tokens": 100, "output_tokens": 30}},
    )

    assert reward_model.score_example(example).score == heuristic.score_example(example).score


def test_learned_preference_scorer_trains_saves_and_loads(tmp_path):
    positive = LearningExample(
        run_id="run-pos",
        trace_id="trace-pos",
        agent_name="bot",
        input_text="Answer from docs",
        output_text="Strong answer",
        created_at="2026-03-18T00:00:00+00:00",
        workspace_id="acme",
        prompt_version="v2",
        model="gpt-4o-mini",
        feedback=[
            _feedback(
                feedback_id="fb-pos", trace_id="trace-pos", run_id="run-pos", kind="thumbs_up"
            )
        ],
        outcomes=[
            _outcome(
                outcome_id="out-pos",
                trace_id="trace-pos",
                run_id="run-pos",
                kind="conversion",
                value=True,
            )
        ],
        metadata={"latency_ms": 150, "usage": {"input_tokens": 400, "output_tokens": 120}},
    )
    negative = LearningExample(
        run_id="run-neg",
        trace_id="trace-neg",
        agent_name="bot",
        input_text="Answer from docs",
        output_text="Bad answer",
        created_at="2026-03-18T00:00:00+00:00",
        workspace_id="acme",
        prompt_version="v2",
        model="claude-sonnet-4-20250514",
        feedback=[
            _feedback(
                feedback_id="fb-neg",
                trace_id="trace-neg",
                run_id="run-neg",
                kind="edited_answer",
                corrected_output="Correct grounded answer",
                expected_output="Correct grounded answer",
                failure_types=["policy_violation"],
            )
        ],
        outcomes=[
            _outcome(
                outcome_id="out-neg",
                trace_id="trace-neg",
                run_id="run-neg",
                kind="human_takeover",
                value=True,
            )
        ],
        failure_types=["policy_violation"],
        corrected_output="Correct grounded answer",
        metadata={"latency_ms": 9000, "usage": {"input_tokens": 50000, "output_tokens": 20000}},
    )

    scorer = LearnedPreferenceScorer.train_from_examples([positive, positive, negative, negative])
    positive_score = scorer.score_example(positive)
    negative_score = scorer.score_example(negative)

    assert positive_score.score > negative_score.score
    assert scorer.training_summary["sample_size"] == 4
    assert scorer.training_summary["preference_pairs"] >= 2

    model_path = tmp_path / "learned_reward.json"
    scorer.save(str(model_path))
    restored = LearnedPreferenceScorer.load(str(model_path))
    assert restored.score_example(positive).score == positive_score.score


def test_training_export_builder_accepts_learned_preference_scorer(tmp_path):
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    learning_store.create_run(
        _run(
            run_id="run-learned-export",
            trace_id="trace-learned-export",
            input_text="Summarize politely",
            output_text="rough answer",
            metadata={"task_type": "support", "latency_ms": 600},
        )
    )
    learning_store.create_feedback(
        _feedback(
            feedback_id="fb-learned-export",
            trace_id="trace-learned-export",
            run_id="run-learned-export",
            kind="edited_answer",
            corrected_output="polite answer",
            expected_output="polite answer",
            failure_types=["style_miss"],
        )
    )
    learning_store.create_outcome(
        _outcome(
            outcome_id="out-learned-export",
            trace_id="trace-learned-export",
            run_id="run-learned-export",
            kind="resolution",
            value=True,
        )
    )

    scorer = LearnedPreferenceScorer.train_from_store(
        learning_store, workspace_id="acme", agent_name="bot"
    )
    builder = TrainingExportBuilder(learning_store, reward_model=scorer)
    rl_export = builder.export_rl(workspace_id="acme", agent_name="bot")

    assert rl_export.records
    assert rl_export.records[0]["reward"] >= 0.0
    assert rl_export.records[0]["components"]


@pytest.mark.skip(reason="API changed - needs update")
def test_policy_learner_combines_model_retrieval_and_tool_histories(tmp_path):
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    routing = ModelRoutingLearningEngine(path=str(tmp_path / "routing.db"))
    retrieval = RetrievalLearningEngine(path=str(tmp_path / "retrieval.db"))
    tool_workflow = ToolWorkflowLearningEngine(path=str(tmp_path / "tool_workflow.db"))

    run = _run(
        run_id="run-1",
        trace_id="trace-1",
        input_text="Find the right policy and answer briefly",
        output_text="Approved answer",
        metadata={
            "task_type": "support",
            "user_segment": "vip",
            "latency_ms": 280,
        },
    )
    learning_store.create_run(run)
    learning_store.create_feedback(
        _feedback(
            feedback_id="fb-1",
            trace_id="trace-1",
            run_id="run-1",
            kind="thumbs_up",
        )
    )
    learning_store.create_outcome(
        _outcome(
            outcome_id="out-1",
            trace_id="trace-1",
            run_id="run-1",
            kind="conversion",
            value=True,
        )
    )

    for _ in range(3):
        routing.record_route(
            model_key="MockProvider:strong",
            success=True,
            latency_ms=40,
            cost_usd=0.002,
            confidence_tier="low",
            workspace_id="acme",
            task_type="support",
            quality_score=0.95,
        )
        routing.record_route(
            model_key="MockProvider:cheap",
            success=False,
            latency_ms=20,
            cost_usd=0.0002,
            confidence_tier="low",
            workspace_id="acme",
            task_type="support",
            quality_score=0.15,
        )

    retrieval.record_retrieval(
        trace_id="trace-1",
        session_id="session-1",
        workspace_id="acme",
        task_type="support",
        user_segment="vip",
        prompt_name="support",
        prompt_version="v1",
        original_query="return policy",
        effective_query="customer return policy approved answer",
        limit=4,
        filter={"category": "policy"},
        chunks=[
            Chunk(
                content="Return policy says answer with approved wording.",
                metadata={"source_key": "docs://policy"},
                document_id="doc-1",
            )
        ],
    )
    retrieval.sync_trace("trace-1", learning_store)

    tool_workflow.record_tool_call(
        tool_name="policy_lookup",
        success=True,
        duration_ms=12,
        task_type="support",
        workspace_id="acme",
    )
    tool_workflow.record_tool_call(
        tool_name="calculator",
        success=False,
        duration_ms=30,
        task_type="support",
        workspace_id="acme",
    )
    tool_workflow.record_fallback(
        provider_key="MockProvider:backup-fast",
        success=True,
        duration_ms=15,
        task_type="support",
        workspace_id="acme",
    )

    learner = PolicyLearner(
        learning_store,
        model_routing=routing,
        retrieval_learning=retrieval,
        tool_workflow=tool_workflow,
    )
    decision = learner.learn_policy(
        workspace_id="acme",
        task_type="support",
        user_segment="vip",
    )

    assert decision.model_key == "MockProvider:strong"
    assert decision.retrieval_limit == 4
    assert decision.retrieval_filter == {"category": "policy"}
    assert decision.preferred_tools[0] == "policy_lookup"
    assert decision.fallback_provider_keys[0] == "MockProvider:backup-fast"
    assert decision.average_reward > 0.5
    assert decision.confidence > 0.5


def test_training_export_builder_outputs_sft_dpo_and_rl_records(tmp_path):
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    learning_store.create_run(
        _run(
            run_id="run-export",
            trace_id="trace-export",
            input_text="Rewrite the answer politely",
            output_text="rough answer",
            metadata={"task_type": "support", "latency_ms": 350},
        )
    )
    learning_store.create_feedback(
        _feedback(
            feedback_id="fb-export",
            trace_id="trace-export",
            run_id="run-export",
            kind="edited_answer",
            corrected_output="polite answer",
            expected_output="polite answer",
            failure_types=["style_miss"],
        )
    )
    learning_store.create_outcome(
        _outcome(
            outcome_id="out-export",
            trace_id="trace-export",
            run_id="run-export",
            kind="resolution",
            value=True,
        )
    )

    builder = TrainingExportBuilder(learning_store)
    sft_export = builder.export_sft(workspace_id="acme", agent_name="bot")
    dpo_export = builder.export_dpo(workspace_id="acme", agent_name="bot")
    rl_export = builder.export_rl(workspace_id="acme", agent_name="bot")

    assert sft_export.kind == "sft"
    assert sft_export.records[0]["completion"] == "polite answer"
    assert dpo_export.kind == "dpo"
    assert dpo_export.records[0]["chosen"] == "polite answer"
    assert dpo_export.records[0]["rejected"] == "rough answer"
    assert rl_export.kind == "rl"
    assert rl_export.records[0]["reward"] > 0.0
    assert '"prompt": "Rewrite the answer politely"' in sft_export.to_jsonl()


@pytest.mark.asyncio
async def test_critic_reviser_loop_returns_improved_answer_and_prompt_patch():
    async def critic(_: str) -> str:
        return json.dumps(
            {
                "should_revise": True,
                "summary": "Answer is too generic",
                "issues": [
                    {
                        "kind": "style_miss",
                        "severity": "medium",
                        "description": "Needs concrete steps",
                    }
                ],
                "prompt_patch": "Require concise numbered steps.",
            }
        )

    def reviser(_: str) -> dict[str, str]:
        return {
            "improved_answer": "1. Check the logs. 2. Restart the worker. 3. Retry the request.",
            "prompt_patch": "Always answer with numbered operational steps.",
            "rationale": "Adds specific executable guidance.",
        }

    loop = CriticReviserLoop(critic=critic, reviser=reviser)
    result = await loop.improve(
        prompt="How do I recover the worker?",
        answer="Try fixing it.",
        failure_types=["style_miss"],
    )

    assert result.revised is True
    assert result.review.should_revise is True
    assert result.review.issues[0].kind == "style_miss"
    assert result.revision.improved_answer.startswith("1.")
    assert "numbered" in (result.revision.prompt_patch or "")


def test_pattern_miner_saves_structured_lessons_to_reflection_memory(tmp_path):
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    reflection_memory = ReflectionMemory(str(tmp_path / "reflection.db"))

    for index in range(2):
        run_id = f"run-{index}"
        trace_id = f"trace-{index}"
        learning_store.create_run(
            _run(
                run_id=run_id,
                trace_id=trace_id,
                input_text="Answer from docs",
                output_text="wrong grounded answer",
                metadata={"task_type": "support", "user_segment": "vip"},
            )
        )
        learning_store.create_feedback(
            _feedback(
                feedback_id=f"fb-{index}",
                trace_id=trace_id,
                run_id=run_id,
                kind="edited_answer",
                corrected_output="correct answer",
                failure_types=["bad_retrieval"],
            )
        )

    miner = PatternMiner(learning_store)
    patterns = miner.mine_patterns(workspace_id="acme", agent_name="bot", min_occurrences=2)
    lessons = miner.save_lessons(
        reflection_memory, workspace_id="acme", agent_name="bot", min_occurrences=2
    )
    persisted = reflection_memory.list_lessons(workspace_id="acme", agent_name="bot")

    assert len(patterns) == 1
    assert patterns[0].signature == "bad_retrieval"
    assert patterns[0].suggested_fix_type == "retrieval_policy"
    assert len(lessons) == 1
    assert persisted[0].category == "retrieval_policy"
    assert persisted[0].metadata["signature"] == "bad_retrieval"


@pytest.mark.skip(reason="API changed - needs update")
def test_root_cause_analyzer_classifies_and_ranks_hypotheses(tmp_path):
    learning_store = SQLiteLearningStore(str(tmp_path / "learning.db"))
    learning_store.create_run(
        _run(
            run_id="run-rca-1",
            trace_id="trace-rca-1",
            input_text="Answer from policy docs",
            output_text="bad answer",
            metadata={
                "task_type": "support",
                "user_segment": "vip",
                "prompt_variant": "baseline",
                "latency_ms": 9200,
                "retrieval_profile": "limit=3|refund-policy",
                "tool_calls": [{"name": "policy_lookup"}],
                "retrieval_source_keys": ["docs://refund-policy"],
            },
            usage={"input_tokens": 15000, "output_tokens": 10000},
        )
    )
    learning_store.create_feedback(
        _feedback(
            feedback_id="fb-rca-1",
            trace_id="trace-rca-1",
            run_id="run-rca-1",
            kind="edited_answer",
            corrected_output="approved answer",
        )
    )
    learning_store.create_outcome(
        _outcome(
            outcome_id="out-rca-1",
            trace_id="trace-rca-1",
            run_id="run-rca-1",
            kind="human_takeover",
            value=True,
        )
    )
    learning_store.create_run(
        _run(
            run_id="run-rca-2",
            trace_id="trace-rca-2",
            input_text="Answer from policy docs",
            output_text="bad answer",
            metadata={
                "task_type": "support",
                "user_segment": "vip",
                "prompt_variant": "baseline",
                "latency_ms": 8700,
                "retrieval_profile": "limit=3|refund-policy",
                "tool_calls": [{"name": "policy_lookup"}],
                "retrieval_source_keys": ["docs://refund-policy"],
            },
            usage={"input_tokens": 16000, "output_tokens": 11000},
        )
    )
    learning_store.create_feedback(
        _feedback(
            feedback_id="fb-rca-2",
            trace_id="trace-rca-2",
            run_id="run-rca-2",
            kind="edited_answer",
            corrected_output="approved answer",
        )
    )

    report = RootCauseAnalyzer(learning_store).analyze(
        workspace_id="acme",
        agent_name="bot",
        limit=50,
        min_occurrences=2,
    )

    assert report.sample_size == 2
    assert report.automatic_classification_count == 2
    assert report.failure_summary["wrong_tool"] >= 1
    assert report.failure_summary["latency"] >= 1
    assert report.failure_summary["too_expensive"] >= 1
    assert report.correlations["prompt_version"][0]["value"] == "baseline"
    assert report.correlations["tool_path"][0]["value"] == "policy_lookup"
    assert report.correlations["retrieval_profile"][0]["value"] == "limit=3|refund-policy"
    assert report.hypotheses
    assert report.hypotheses[0].confidence > 0.0
    assert report.hypotheses[0].affected_surface in {
        "prompt",
        "tool_policy",
        "retrieval_policy",
        "routing_policy",
        "workflow_policy",
        "guardrail",
    }
