"""
Live integration test for mubit-agno against a running MuBit instance.

Requires: MuBit running at http://127.0.0.1:3000 (make run-mubit)
"""

from __future__ import annotations

import json
import os
import time
import uuid

from mubit_agno import MubitMemoryDb, MubitToolkit, MubitAgnoMemory


SESSION = f"agno-live-test-{uuid.uuid4().hex[:8]}"
USER = "test-user"
ENDPOINT = "http://127.0.0.1:3000"
API_KEY = os.environ.get(
    "MUBIT_API_KEY",
    "mbt_local_f56d24f60b1f4346_317f8d47d3264c1aad3f7e0ef560e0c06841977e245c44ef8cc611a052d15d6f",
)

# Shared client so all tests use the same session
_client = None


def _get_client():
    global _client
    if _client is None:
        from mubit_agno.mubit_client import MubitControlClient
        _client = MubitControlClient(endpoint=ENDPOINT, api_key=API_KEY)
    return _client


def _ingest_seed_data():
    """Seed the session with data and wait for indexing."""
    client = _get_client()
    client.remember(
        session_id=SESSION,
        user_id=USER,
        text="The project uses PostgreSQL 16 on AWS RDS in us-east-1",
        agent_id="seed",
        intent="fact",
        upsert_key="seed-pg",
        metadata={"agno_memory_id": "seed-pg", "topics": ["database", "infrastructure"]},
    )
    client.remember(
        session_id=SESSION,
        user_id=USER,
        text="Redis is used as the caching layer on port 6379",
        agent_id="seed",
        intent="fact",
        upsert_key="seed-redis",
        metadata={"agno_memory_id": "seed-redis", "topics": ["cache", "infrastructure"]},
    )
    client.remember(
        session_id=SESSION,
        user_id=USER,
        text="Always validate user input before database queries to prevent SQL injection",
        agent_id="seed",
        intent="lesson",
        upsert_key="seed-lesson-1",
        lesson_type="rule",
        lesson_scope="global",
        metadata={"agno_memory_id": "seed-lesson-1", "topics": ["security"]},
    )
    # Wait for indexing
    time.sleep(2)


def test_00_seed():
    """Seed test data into MuBit."""
    _ingest_seed_data()
    print("  PASS: seed data ingested")


def test_memory_db_read():
    """MubitMemoryDb: read memories back from seeded data."""
    db = MubitMemoryDb(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-db",
        mubit_client=_get_client(),
    )
    assert db.table_exists()

    rows = db.read_memories(user_id=USER, limit=10)
    print(f"  read_memories returned {len(rows)} rows")
    assert len(rows) >= 1
    contents = [getattr(r, "memory", "") for r in rows]
    print(f"  contents: {contents[:3]}")
    assert any("PostgreSQL" in c or "Redis" in c for c in contents), f"Expected DB info in {contents}"
    print("  PASS: read_memories works")


def test_memory_db_upsert():
    """MubitMemoryDb: upsert a new memory."""
    db = MubitMemoryDb(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-db",
        mubit_client=_get_client(),
    )

    class FakeMemory:
        id = "upsert-mem-1"
        memory = "The API gateway uses Kong with rate limiting enabled"
        user_id = USER
        topics = ["api", "infrastructure"]
        memory_type = "semantic"

    result = db.upsert_memory(FakeMemory())
    print(f"  upsert_memory returned: {result}")
    time.sleep(1)

    # Verify it's retrievable
    rows = db.read_memories(user_id=USER, limit=20)
    contents = [getattr(r, "memory", "") for r in rows]
    assert any("Kong" in c for c in contents), f"Expected Kong in {contents}"
    print("  PASS: upsert_memory works")


def test_memory_db_delete():
    """MubitMemoryDb: delete marks memory as deleted."""
    db = MubitMemoryDb(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-db",
        mubit_client=_get_client(),
    )
    db.delete_memory("upsert-mem-1")
    time.sleep(0.3)
    print("  PASS: delete_memory did not error")


def test_toolkit_recall():
    """MubitToolkit: recall seeded data."""
    tk = MubitToolkit(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-toolkit",
        mubit_client=_get_client(),
    )

    result = tk.mubit_recall("What database technology do we use?", limit=5)
    parsed = json.loads(result)
    evidence = parsed.get("evidence", [])
    print(f"  recall returned {len(evidence)} evidence items")
    assert len(evidence) >= 1
    contents = [e.get("content", "") for e in evidence]
    print(f"  top results: {[c[:60] for c in contents[:3]]}")
    assert any("PostgreSQL" in c or "Redis" in c for c in contents), f"Expected DB info in {contents}"
    print("  PASS: toolkit recall works")


def test_toolkit_remember():
    """MubitToolkit: remember stores new data."""
    tk = MubitToolkit(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-toolkit",
        mubit_client=_get_client(),
    )

    result = tk.mubit_remember(
        "Deployments use ArgoCD with GitOps workflow",
        intent="fact",
    )
    parsed = json.loads(result)
    print(f"  remember result: job_id={parsed.get('job_id')}")
    assert "job_id" in parsed
    print("  PASS: toolkit remember works")


def test_toolkit_get_context():
    """MubitToolkit: get_context returns assembled context."""
    tk = MubitToolkit(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-toolkit",
        mubit_client=_get_client(),
    )

    result = tk.mubit_get_context("infrastructure overview")
    parsed = json.loads(result)
    print(f"  get_context keys: {list(parsed.keys())}")
    assert isinstance(parsed, dict)
    print("  PASS: get_context works")


def test_toolkit_checkpoint():
    """MubitToolkit: checkpoint creates a snapshot."""
    tk = MubitToolkit(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-toolkit",
        mubit_client=_get_client(),
    )

    result = tk.mubit_checkpoint("live-test-phase-1", "Completed seed and recall tests")
    parsed = json.loads(result)
    print(f"  checkpoint result: {parsed}")
    assert isinstance(parsed, dict)
    print("  PASS: checkpoint works")


def test_toolkit_memory_health():
    """MubitToolkit: memory_health returns status."""
    tk = MubitToolkit(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-toolkit",
        mubit_client=_get_client(),
    )

    result = tk.mubit_memory_health()
    parsed = json.loads(result)
    print(f"  memory_health keys: {list(parsed.keys())}")
    assert isinstance(parsed, dict)
    print("  PASS: memory_health works")


def test_convenience_register_and_list_agents():
    """MubitAgnoMemory: register + list agents."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    result = mem.register_agent(
        "researcher",
        role="researcher",
        read_scopes=["fact", "lesson"],
        write_scopes=["fact", "trace"],
        shared_memory_lanes=["knowledge"],
    )
    print(f"  register_agent: {result.get('success')}")
    assert result.get("success") is True

    result = mem.register_agent(
        "writer",
        role="writer",
        read_scopes=["fact", "lesson", "trace"],
        write_scopes=["trace"],
    )
    assert result.get("success") is True

    agents = mem.list_agents()
    agent_ids = [a["agent_id"] for a in agents.get("agents", [])]
    print(f"  agents: {agent_ids}")
    assert "researcher" in agent_ids
    assert "writer" in agent_ids
    print("  PASS: register + list agents works")


def test_convenience_handoff():
    """MubitAgnoMemory: create a handoff between agents."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    result = mem.handoff(
        "researcher", "writer",
        "Research complete. Key findings stored in memory.",
        requested_action="review",
    )
    print(f"  handoff result: {result}")
    assert isinstance(result, dict)
    print("  PASS: handoff works")


def test_convenience_get_context():
    """MubitAgnoMemory: get assembled context."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    ctx = mem.get_context("database infrastructure")
    print(f"  get_context length: {len(ctx)} chars")
    print(f"  context preview: {ctx[:200]}...")
    print("  PASS: get_context works")


def test_convenience_record_outcome():
    """MubitAgnoMemory: record outcome."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    # First recall to get a valid evidence ID
    data = _get_client().recall(
        session_id=SESSION,
        user_id=USER,
        query="PostgreSQL database",
        limit=1,
    )
    evidence = data.get("evidence", [])
    ref_id = evidence[0]["id"] if evidence else "seed-pg"

    result = mem.record_outcome(
        ref_id,
        "success",
        signal=1.0,
        rationale="Successfully retrieved seeded infrastructure data",
    )
    print(f"  record_outcome: {result}")
    assert isinstance(result, dict)
    print("  PASS: record_outcome works")


def test_convenience_surface_strategies():
    """MubitAgnoMemory: surface strategies."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    strategies = mem.surface_strategies(max_strategies=3)
    print(f"  strategies: {strategies}")
    assert isinstance(strategies, dict)
    print("  PASS: surface_strategies works")


def test_convenience_reflect():
    """MubitAgnoMemory: reflect extracts lessons."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    result = mem.reflect()
    print(f"  reflect: lessons_stored={result.get('lessons_stored')}, summary={result.get('summary', '')[:80]}")
    assert isinstance(result, dict)
    print("  PASS: reflect works")


def test_convenience_lessons():
    """MubitAgnoMemory: list lessons."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    result = mem.lessons(limit=10)
    print(f"  lessons: {result}")
    assert isinstance(result, dict)
    print("  PASS: lessons works")


def test_convenience_diagnose():
    """MubitAgnoMemory: diagnose error patterns."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    result = mem.diagnose("connection timeout to database")
    print(f"  diagnose: {result}")
    assert isinstance(result, dict)
    print("  PASS: diagnose works")


def test_convenience_archive_dereference():
    """MubitAgnoMemory: archive + dereference round-trip."""
    mem = MubitAgnoMemory(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-wrapper",
        mubit_client=_get_client(),
    )

    try:
        archived = mem.archive(
            "SELECT * FROM users WHERE active = true",
            "sql_query",
            labels=["database", "users"],
            family="queries",
        )
        print(f"  archive result: {archived}")
        ref_id = archived.get("reference_id")
        if ref_id:
            deref = mem.dereference(ref_id)
            print(f"  dereference result: {deref}")
            actual_content = deref.get("content") or deref.get("evidence", {}).get("content")
            assert actual_content == "SELECT * FROM users WHERE active = true", f"Got: {actual_content}"
            print("  PASS: archive + dereference round-trip works")
        else:
            print("  PASS (partial): archive succeeded but no reference_id")
    except Exception as e:
        if "embedding" in str(e).lower() or "encode" in str(e).lower():
            print(f"  SKIP: embedding service not available ({type(e).__name__})")
        else:
            raise


def test_zz_cleanup():
    """Clean up test data."""
    db = MubitMemoryDb(
        endpoint=ENDPOINT,
        api_key=API_KEY,
        session_id=SESSION,
        user_id=USER,
        agent_id="test-cleanup",
        mubit_client=_get_client(),
    )
    db.clear()
    print("  PASS: cleanup done")


if __name__ == "__main__":
    tests = [
        ("Seed data", test_00_seed),
        ("MemoryDb read", test_memory_db_read),
        ("MemoryDb upsert", test_memory_db_upsert),
        ("MemoryDb delete", test_memory_db_delete),
        ("Toolkit recall", test_toolkit_recall),
        ("Toolkit remember", test_toolkit_remember),
        ("Toolkit get_context", test_toolkit_get_context),
        ("Toolkit checkpoint", test_toolkit_checkpoint),
        ("Toolkit memory_health", test_toolkit_memory_health),
        ("Register + list agents", test_convenience_register_and_list_agents),
        ("Handoff", test_convenience_handoff),
        ("Get context", test_convenience_get_context),
        ("Record outcome", test_convenience_record_outcome),
        ("Surface strategies", test_convenience_surface_strategies),
        ("Reflect", test_convenience_reflect),
        ("Lessons", test_convenience_lessons),
        ("Diagnose", test_convenience_diagnose),
        ("Archive + dereference", test_convenience_archive_dereference),
        ("Cleanup", test_zz_cleanup),
    ]

    passed = 0
    failed = 0
    for name, fn in tests:
        print(f"\n[TEST] {name}")
        try:
            fn()
            passed += 1
        except Exception as e:
            print(f"  FAIL: {e}")
            failed += 1

    print(f"\n{'='*50}")
    print(f"Results: {passed} passed, {failed} failed out of {len(tests)}")
    if failed:
        exit(1)
