#!/usr/bin/env python3
"""
Real-world functional tests for Learning Loop.
Tests evaluation, feedback capture, pattern mining, reflection memory, and improvement.
"""

import asyncio
import json
import os
import sys
import tempfile
from pathlib import Path
from typing import Any
from uuid import uuid4

sys.path.insert(0, os.path.dirname(__file__))


def get_provider():
    return dict(
        model=os.getenv("OPENAI_COMPAT_MODEL", "gpt-4o-mini"),
        base_url=os.getenv("OPENAI_COMPAT_BASE_URL", "https://api.example.com/v1"),
        api_key=os.getenv("OPENAI_COMPAT_API_KEY", "test-api-key"),
    )


# ============================================================================
# TEST 1: Feedback Capture - Store evaluation results
# ============================================================================


def test_feedback_capture_sqlite():
    """Test: Capture feedback/evaluation results to learning store."""
    print("\n" + "=" * 60)
    print("TEST 1: Feedback Capture (SQLite)")
    print("=" * 60)

    from logos_adk.learning.store import SQLiteLearningStore
    from logos_adk.learning.models import FeedbackRecord

    with tempfile.TemporaryDirectory() as tmpdir:
        store = SQLiteLearningStore(path=f"{tmpdir}/learning.db")

        feedback1 = FeedbackRecord(
            id=str(uuid4()),
            kind="edited_answer",
            trace_id="trace-1",
            run_id="run-1",
            corrected_output="42",
            failure_types=["hallucination"],
            workspace_id="test-ws",
        )

        feedback2 = FeedbackRecord(
            id=str(uuid4()),
            kind="numeric_rating",
            trace_id="trace-2",
            run_id="run-2",
            rating=0.8,
            workspace_id="test-ws",
        )

        store.create_feedback(feedback1)
        store.create_feedback(feedback2)

        all_fb = store.list_feedback(workspace_id="test-ws")
        print(f"  Stored: {len(all_fb)} feedback records")

        assert len(all_fb) == 2

    print("✓ TEST 1 PASSED\n")
    return True


def test_feedback_capture_postgresql():
    """Test: Capture feedback to PostgreSQL."""
    print("\n" + "=" * 60)
    print("TEST 2: Feedback Capture (PostgreSQL)")
    print("=" * 60)

    dsn = os.getenv("LOGOS_ADK_DATABASE_URL", "postgresql://crm:crm@localhost:5432/crm")

    try:
        from logos_adk.learning.store import PostgreSQLLearningStore
        from logos_adk.learning.models import FeedbackRecord

        table = f"test_fb_{uuid4().hex[:8]}"
        store = PostgreSQLLearningStore(dsn=dsn, table_prefix=table)

        fb = FeedbackRecord(
            id=str(uuid4()),
            kind="thumbs_down",
            trace_id="pg-trace-1",
            failure_types=["hallucination"],
            workspace_id="pg-ws",
        )

        store.create_feedback(fb)
        print(f"  Stored feedback in PostgreSQL: {fb.id}")

    except Exception as e:
        print(f"  PostgreSQL test skipped: {e}")
        return True

    print("✓ TEST 2 PASSED\n")
    return True


# ============================================================================
# TEST 3: Run Records
# ============================================================================


def test_run_records():
    """Test: Store and retrieve run records."""
    print("\n" + "=" * 60)
    print("TEST 3: Run Records Storage")
    print("=" * 60)

    from logos_adk.learning.store import SQLiteLearningStore
    from logos_adk.learning.models import RunRecord

    with tempfile.TemporaryDirectory() as tmpdir:
        store = SQLiteLearningStore(path=f"{tmpdir}/learning.db")

        run = RunRecord(
            id=str(uuid4()),
            trace_id="trace-run-1",
            request_id="req-1",
            agent_name="assistant",
            input_text="What is 2+2?",
            output_text="4",
        )

        store.create_run(run)

        retrieved = store.get_run(run.id)
        assert retrieved is not None
        assert retrieved.input_text == "What is 2+2?"

        print(f"  Stored run: {run.id}")

    print("✓ TEST 3 PASSED\n")
    return True


# ============================================================================
# TEST 4: Reflection Memory
# ============================================================================


def test_reflection_memory_sqlite():
    """Test: Reflection memory stores lessons."""
    print("\n" + "=" * 60)
    print("TEST 4: Reflection Memory (SQLite)")
    print("=" * 60)

    from logos_adk.learning.improvement.reflection_memory import ReflectionMemory
    from logos_adk.learning.improvement.models import ReflectionLesson

    with tempfile.TemporaryDirectory() as tmpdir:
        memory = ReflectionMemory(path=f"{tmpdir}/reflection.db")

        lesson = ReflectionLesson(
            id=str(uuid4()),
            workspace_id="test-ws",
            agent_name="assistant",
            category="incorrect_reasoning",
            lesson="Verify math calculations",
            rationale="Multiple feedback records showed errors",
            suggested_action="Add verification step",
            confidence=0.85,
            source_run_ids=["run-1"],
            metadata={"failure_count": 5},
        )

        memory.save_lesson(lesson)
        print(f"  Stored lesson: {lesson.id}")

        lessons = memory.list_lessons(workspace_id="test-ws")
        print(f"  Retrieved: {len(lessons)} lessons")

        assert len(lessons) >= 1

    print("✓ TEST 4 PASSED\n")
    return True


def test_reflection_memory_postgresql():
    """Test: Reflection memory with PostgreSQL."""
    print("\n" + "=" * 60)
    print("TEST 5: Reflection Memory (PostgreSQL)")
    print("=" * 60)

    dsn = os.getenv("LOGOS_ADK_DATABASE_URL", "postgresql://crm:crm@localhost:5432/crm")

    try:
        from logos_adk.learning.improvement.reflection_memory import ReflectionMemory
        from logos_adk.learning.improvement.models import ReflectionLesson

        table = f"test_ref_{uuid4().hex[:8]}"
        memory = ReflectionMemory(dsn=dsn, table_name=table)

        lesson = ReflectionLesson(
            id=str(uuid4()),
            workspace_id="pg-ws",
            agent_name="pg_agent",
            category="test_cat",
            lesson="Test lesson",
            rationale="Testing",
            suggested_action="Verify",
            confidence=0.9,
            source_run_ids=["pg-run-1"],
            metadata={"test": True},
        )

        memory.save_lesson(lesson)
        print(f"  Stored in PostgreSQL: {lesson.id}")

    except Exception as e:
        print(f"  PostgreSQL test skipped: {e}")
        return True

    print("✓ TEST 5 PASSED\n")
    return True


# ============================================================================
# TEST 6: Pattern Mining
# ============================================================================


def test_pattern_mining():
    """Test: Mine patterns from feedback."""
    print("\n" + "=" * 60)
    print("TEST 6: Pattern Mining")
    print("=" * 60)

    from logos_adk.learning.store import SQLiteLearningStore
    from logos_adk.learning.improvement.pattern_miner import PatternMiner
    from logos_adk.learning.models import FeedbackRecord

    with tempfile.TemporaryDirectory() as tmpdir:
        store = SQLiteLearningStore(path=f"{tmpdir}/learning.db")

        for i in range(5):
            fb = FeedbackRecord(
                id=str(uuid4()),
                kind="edited_answer",
                trace_id=f"trace-{i}",
                run_id=f"run-{i}",
                corrected_output=f"Correct: {i * 10}",
                failure_types=["hallucination"],
                workspace_id="pattern-ws",
            )
            store.create_feedback(fb)

        print(f"  Stored 5 feedback records")

        miner = PatternMiner(store)
        patterns = miner.mine_patterns(
            workspace_id="pattern-ws",
            agent_name="assistant",
            limit=10,
            min_occurrences=2,
        )

        print(f"  Mined patterns: {len(patterns)}")

    print("✓ TEST 6 PASSED\n")
    return True


# ============================================================================
# TEST 7: Full Learning Loop
# ============================================================================


async def test_full_learning_loop():
    """Test: Complete learning loop."""
    print("\n" + "=" * 60)
    print("TEST 7: Full Learning Loop")
    print("=" * 60)

    from logos_adk.learning.store import SQLiteLearningStore
    from logos_adk.learning.improvement.reflection_memory import ReflectionMemory
    from logos_adk.learning.improvement.models import ReflectionLesson
    from logos_adk.learning.models import FeedbackRecord, RunRecord
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    with tempfile.TemporaryDirectory() as tmpdir:
        store = SQLiteLearningStore(path=f"{tmpdir}/learning.db")
        reflection = ReflectionMemory(path=f"{tmpdir}/reflection.db")

        print("  Step 1: Setup stores ✓")

        provider = OpenAICompatibleProvider(**get_provider())
        agent = Agent("learning_agent", model=provider)

        prompt = "What is 15 + 27?"
        response = await agent.run(prompt)
        output = response.content

        print(f"  Step 2: Agent executed: {output[:30]}...")

        run = RunRecord(
            id=str(uuid4()),
            trace_id=str(uuid4()),
            request_id=str(uuid4()),
            agent_name="learning_agent",
            input_text=prompt,
            output_text=output,
        )
        store.create_run(run)

        print(f"  Step 3: Stored run record ✓")

        fb = FeedbackRecord(
            id=str(uuid4()),
            kind="edited_answer",
            trace_id=run.trace_id,
            run_id=run.id,
            corrected_output="Expected 42",
            failure_types=["hallucination"],
            workspace_id="loop-ws",
        )
        store.create_feedback(fb)

        print(f"  Step 4: Stored feedback ✓")

        lesson = ReflectionLesson(
            id=str(uuid4()),
            workspace_id="loop-ws",
            agent_name="learning_agent",
            category="math_accuracy",
            lesson="Double-check arithmetic",
            rationale="Feedback showed math errors",
            suggested_action="Add verification",
            confidence=0.75,
            source_run_ids=[run.id],
            metadata={"prompt_type": "math"},
        )
        reflection.save_lesson(lesson)

        print(f"  Step 5: Stored lesson ✓")

        lessons = reflection.list_lessons(workspace_id="loop-ws")

        print(f"  Step 6: Retrieved {len(lessons)} lessons")

        response2 = await agent.run(
            f"{prompt}. Remember: {lessons[0].lesson if lessons else 'none'}"
        )

        print(f"  Step 7: Second run with learned context")
        print(f"    Output: {response2.content[:30]}...")

        final = reflection.list_lessons(workspace_id="loop-ws")

        print(f"\n  === LEARNING SUMMARY ===")
        print(f"  Lessons stored: {len(final)}")
        print(f"  Lessons applied: Yes")

    print("✓ TEST 7 PASSED\n")
    return True


# ============================================================================
# TEST 8: Dataset Builder
# ============================================================================


def test_dataset_builder():
    """Test: Build training datasets."""
    print("\n" + "=" * 60)
    print("TEST 8: Dataset Builder")
    print("=" * 60)

    from logos_adk.learning.store import SQLiteLearningStore
    from logos_adk.learning.dataset import DatasetBuilder
    from logos_adk.learning.models import FeedbackRecord, RunRecord

    with tempfile.TemporaryDirectory() as tmpdir:
        store = SQLiteLearningStore(path=f"{tmpdir}/learning.db")

        for i in range(3):
            run = RunRecord(
                id=f"run-ds-{i}",
                trace_id=f"trace-ds-{i}",
                request_id=f"req-ds-{i}",
                agent_name="assistant",
                input_text=f"What is {i} * 10?",
                output_text=f"Wrong: {i * 100}",
            )
            store.create_run(run)

            fb = FeedbackRecord(
                id=str(uuid4()),
                kind="edited_answer",
                trace_id=f"trace-ds-{i}",
                run_id=run.id,
                corrected_output=f"Correct: {i * 10}",
                failure_types=["hallucination"],
                workspace_id="ds-ws",
            )
            store.create_feedback(fb)

        print(f"  Stored 3 runs with corrections")

        builder = DatasetBuilder(store)
        dataset = builder.build(
            workspace_id="ds-ws",
            agent_name="assistant",
            limit=10,
            include_unreviewed=False,
        )

        print(f"  Built dataset: {len(dataset)} examples")

        for ex in dataset[:2]:
            print(f"    Input: {ex.input_text[:20]}...")
            print(f"    Output: {ex.output_text[:20]}...")
            print(f"    Corrected: {ex.corrected_output[:20] if ex.corrected_output else 'N/A'}...")

    print("✓ TEST 8 PASSED\n")
    return True


# ============================================================================
# MAIN
# ============================================================================


async def main():
    print("=" * 60)
    print("LEARNING LOOP - REAL API & DATABASE TESTS")
    print("=" * 60)
    print(f"API: {get_provider()['base_url']}")
    print(f"Model: {get_provider()['model']}")

    results = []

    sync_tests = [
        ("Feedback (SQLite)", test_feedback_capture_sqlite),
        ("Feedback (PostgreSQL)", test_feedback_capture_postgresql),
        ("Run Records", test_run_records),
        ("Reflection Memory (SQLite)", test_reflection_memory_sqlite),
        ("Reflection Memory (PG)", test_reflection_memory_postgresql),
        ("Pattern Mining", test_pattern_mining),
        ("Dataset Builder", test_dataset_builder),
    ]

    for name, test_func in sync_tests:
        try:
            result = test_func()
            results.append((name, result, None))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False, str(e)))

    async_tests = [
        ("Full Learning Loop", test_full_learning_loop),
    ]

    for name, test_func in async_tests:
        try:
            result = await test_func()
            results.append((name, result, None))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False, str(e)))

    print("=" * 60)
    print("RESULTS")
    print("=" * 60)

    passed = sum(1 for _, ok, _ in results if ok)
    total = len(results)

    for name, ok, err in results:
        status = "✓ PASS" if ok else "✗ FAIL"
        print(f"{status}: {name}")
        if err:
            print(f"    Error: {err}")

    print(f"\n{passed}/{total} tests passed")
    return passed == total


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
