#!/usr/bin/env python3
"""Integration test — verify ALL LLM providers work inside the full neo-cortex system.

Each provider is tested end-to-end:
  1. create_client(provider) → LLMClient
  2. Classifier(llm) → classify a real conversation turn
  3. CortexService.ingest() → full pipeline (classify + embed + store)
  4. CortexService.recall() → verify the ingested memory is retrievable

This catches: wrong API format, bad auth, broken JSON parsing, schema
mismatches, timeout issues — anything that would break production.

Usage:
    cd ~/neo-ram/neo-cortex
    CORTEX_DATA_DIR=~/neo-ram uv run python tests/integration/test_providers.py
    CORTEX_DATA_DIR=~/neo-ram uv run python tests/integration/test_providers.py --only gemini,groq
    CORTEX_DATA_DIR=~/neo-ram uv run python tests/integration/test_providers.py --exclude claude-code,x.ai
"""

import argparse
import asyncio
import json
import os
import shutil
import sys
import tempfile
import time

# Use development code
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))

# All providers with their required env var and default model
PROVIDERS = [
    ("gemini",      "GEMINI_API_KEY",    "gemini-2.5-flash"),
    ("groq",        "GROQ_API_KEY",      "llama-3.3-70b-versatile"),
    ("anthropic",   "ANTHROPIC_API_KEY", "claude-haiku-4-5-20251001"),
    ("openai",      "OPENAI_API_KEY",    "gpt-4.1-nano"),
    ("z.ai",        "ZAI_API_KEY",       "glm-4.7"),
    ("claude-code", None,                "haiku"),
]

# Real conversation sample for testing
SAMPLE_Q = "Come funziona il sistema di ingest delle memorie in neo-cortex?"
SAMPLE_A = (
    "Il pipeline di ingest segue 6 step: subscriber ascolta su SignalR, "
    "TurnAccumulator raccoglie Q&A, Digestor processa i turni, noise filter "
    "scarta contenuto triviale, Groq classifica (project/topic/title/concepts), "
    "Jina embedda a 1024-dim vectors, SimHash deduplica, ChromaDB salva i vettori "
    "e SQLite salva metadata per FTS5."
)


def load_keys():
    """Load API keys from ~/neo-ram/.keys into environment."""
    keys_file = os.path.expanduser("~/neo-ram/.keys")
    if not os.path.exists(keys_file):
        return
    mapping = {
        "gemini": "GEMINI_API_KEY",
        "groq": "GROQ_API_KEY",
        "open-ai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "z.ai": "ZAI_API_KEY",
        "x.ai": "XAI_API_KEY",
        "jina": "JINA_API_KEY",
    }
    with open(keys_file) as f:
        for line in f:
            if ":" in line:
                name, key = line.strip().split(":", 1)
                env_var = mapping.get(name.strip())
                if env_var and env_var not in os.environ:
                    os.environ[env_var] = key.strip()


class ProviderIntegrationTest:
    """Test a single provider through the full neo-cortex pipeline."""

    def __init__(self, provider: str, model: str):
        self.provider = provider
        self.model = model
        self.tmpdir = tempfile.mkdtemp(prefix=f"cortex_prov_{provider}_")
        self.results: list[tuple[str, bool, str]] = []

    def record(self, name: str, passed: bool, detail: str = ""):
        status = "PASS" if passed else "FAIL"
        self.results.append((name, passed, detail))
        color = "\033[32m" if passed else "\033[31m"
        print(f"    [{color}{status}\033[0m] {name}")
        if detail:
            print(f"           {detail}")

    async def run(self) -> bool:
        """Run full pipeline test for this provider. Returns True if all passed."""
        from neo_cortex.classifier import Classifier
        from neo_cortex.cortex import CortexService
        from neo_cortex.dedup import DedupStore
        from neo_cortex.embedder import JinaEmbedder
        from neo_cortex.graph import ConceptGraph
        from neo_cortex.llm import create_client
        from neo_cortex.memory_index import MemoryIndex
        from neo_cortex.models import IngestRequest
        from neo_cortex.store import MemoryStore

        print(f"\n  --- {self.provider}/{self.model} ---")

        # Step 1: Create LLM client
        try:
            llm = create_client(self.provider)
            self.record("create_client", True, f"name={llm.name}")
        except Exception as e:
            self.record("create_client", False, str(e)[:200])
            return False

        # Step 2: Classify a real conversation turn
        try:
            classifier = Classifier(llm)
            t0 = time.monotonic()
            result = await classifier.classify(SAMPLE_Q, SAMPLE_A)
            ms = (time.monotonic() - t0) * 1000

            has_project = bool(result.project)
            has_topic = bool(result.topic)
            has_concepts = len(result.concepts) > 0

            self.record(
                "classify",
                has_project and has_topic,
                f"{ms:.0f}ms project={result.project} topic={result.topic} "
                f"concepts={result.concepts[:3]}",
            )
        except Exception as e:
            self.record("classify", False, str(e)[:200])
            await llm.close()
            return False

        # Step 3: Full ingest through CortexService
        try:
            db_path = os.path.join(self.tmpdir, "chroma")
            index_path = os.path.join(self.tmpdir, "index.db")

            store = MemoryStore(db_path, "provider_test")
            embedder = JinaEmbedder(os.environ["JINA_API_KEY"])
            index = MemoryIndex(index_path)
            index.register_runtime_day()
            dedup = DedupStore(index.conn)
            graph = ConceptGraph(index.conn)

            cortex = CortexService(
                store, embedder, classifier, index=index, dedup=dedup, graph=graph,
            )

            t0 = time.monotonic()
            req = IngestRequest(
                session_id="provider-test",
                question=SAMPLE_Q,
                answer=SAMPLE_A,
                model="test",
                turn_number=0,
            )
            memory_id = await cortex.ingest(req)
            ms = (time.monotonic() - t0) * 1000

            self.record(
                "ingest",
                memory_id is not None,
                f"{ms:.0f}ms memory_id={memory_id}",
            )
        except Exception as e:
            self.record("ingest", False, str(e)[:200])
            await llm.close()
            return False

        # Step 4: Recall — verify the memory is retrievable
        try:
            t0 = time.monotonic()
            recall_result = await cortex.recall("ingest pipeline neo-cortex", n=3, smart=False)
            ms = (time.monotonic() - t0) * 1000

            found = recall_result.count > 0
            self.record(
                "recall",
                found,
                f"{ms:.0f}ms count={recall_result.count}"
                + (f" top_sim={recall_result.memories[0].similarity:.3f}" if found else ""),
            )
        except Exception as e:
            self.record("recall", False, str(e)[:200])

        # Cleanup
        await llm.close()
        return all(p for _, p, _ in self.results)

    def cleanup(self):
        if os.path.exists(self.tmpdir):
            shutil.rmtree(self.tmpdir, ignore_errors=True)


def parse_args():
    parser = argparse.ArgumentParser(description="Integration test for all LLM providers")
    parser.add_argument("--only", help="Test only these providers (comma-separated)")
    parser.add_argument("--exclude", help="Skip these providers (comma-separated)")
    return parser.parse_args()


async def main():
    args = parse_args()
    only = set(args.only.split(",")) if args.only else None
    exclude = set(args.exclude.split(",")) if args.exclude else set()

    load_keys()

    if "JINA_API_KEY" not in os.environ:
        print("ERROR: JINA_API_KEY not set. Run from ~/neo-ram or export manually.")
        sys.exit(1)

    if "CORTEX_DATA_DIR" not in os.environ:
        os.environ["CORTEX_DATA_DIR"] = os.path.expanduser("~/neo-ram")

    providers = [
        (p, model) for p, key_env, model in PROVIDERS
        if (only is None or p in only) and p not in exclude
        and (key_env is None or key_env in os.environ)
    ]

    skipped = [
        p for p, key_env, _ in PROVIDERS
        if (only is None or p in only) and p not in exclude
        and key_env is not None and key_env not in os.environ
    ]

    print("=" * 60)
    print("  NEO-CORTEX PROVIDER INTEGRATION TEST")
    print("  Full pipeline: create_client → classify → ingest → recall")
    print(f"  Providers: {len(providers)} to test" +
          (f", {len(skipped)} skipped (no key)" if skipped else ""))
    print("=" * 60)

    if skipped:
        for p in skipped:
            print(f"  [SKIP] {p} — API key not found")

    all_results = []
    for provider, model in providers:
        test = ProviderIntegrationTest(provider, model)
        try:
            passed = await test.run()
            all_results.append((provider, passed, test.results))
        except Exception as e:
            print(f"    [FATAL] {e}")
            all_results.append((provider, False, [("fatal", False, str(e))]))
        finally:
            test.cleanup()

    # Summary
    print("\n" + "=" * 60)
    print("  SUMMARY")
    print("=" * 60)

    total_pass = 0
    total_fail = 0
    for provider, passed, results in all_results:
        p = sum(1 for _, ok, _ in results if ok)
        f = sum(1 for _, ok, _ in results if not ok)
        total_pass += p
        total_fail += f
        color = "\033[32m" if passed else "\033[31m"
        status = "ALL PASS" if passed else f"{f} FAILED"
        print(f"  {color}{provider:20s} {status}\033[0m ({p}/{p+f} tests)")

    print()
    if total_fail == 0:
        print(f"  \033[32mALL PROVIDERS PASSED: {total_pass} tests\033[0m")
    else:
        print(f"  \033[31m{total_fail} FAILED\033[0m across {len(all_results)} providers")

    print("=" * 60)
    sys.exit(0 if total_fail == 0 else 1)


if __name__ == "__main__":
    asyncio.run(main())
