"""
Test all TraceQuery API methods against the cached result.
No embedding calls needed — uses existing cached data.

Run: python test/test_api_surface.py
"""

import sys
import os
import json
import traceback

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

from tracescope.models.analysis import AnalysisResult
from tracescope.query import TraceQuery
from tracescope.providers.embedding import OpenAIEmbedding

CACHE_PATH = os.path.join(os.path.dirname(__file__), "..", "cache", "prm_demo_40paths_flow.json")

passed = 0
failed = 0
errors = []


def test(name, fn):
    global passed, failed
    try:
        fn()
        print(f"  PASS: {name}")
        passed += 1
    except Exception as e:
        print(f"  FAIL: {name} — {e}")
        traceback.print_exc()
        errors.append((name, str(e)))
        failed += 1


def main():
    global passed, failed

    print("=" * 60)
    print("TraceScope API Surface Test")
    print("=" * 60)

    # ── Load cached result ──
    print("\n--- Loading cached result ---")
    result = AnalysisResult.load_result(CACHE_PATH)
    print(f"Loaded: {result.n_entries} entries, {result.clusters.n_clusters} clusters")

    # ── Test AnalysisResult methods ──
    print("\n--- AnalysisResult methods ---")

    test("fingerprint()", lambda: (
        result.fingerprint(),
        assert_(len(result.fingerprint()) == 64, "fingerprint should be 64 hex chars")
    ))

    test("get_cluster_texts(0)", lambda: (
        texts := result.get_cluster_texts(0),
        assert_(len(texts) > 0, "cluster 0 should have texts"),
        assert_(isinstance(texts[0], str), "texts should be strings")
    ))

    test("get_entry_at_3d(0)", lambda: (
        entry := result.get_entry_at_3d(0),
        assert_("text" in entry, "should have text"),
        assert_("x" in entry, "should have x"),
        assert_("cluster" in entry, "should have cluster")
    ))

    test("n_entries property", lambda: (
        assert_(result.n_entries == 348, f"expected 348, got {result.n_entries}")
    ))

    # ── Test TraceQuery without embedding provider ──
    # (can't embed new text, but can test internal methods)
    print("\n--- TraceQuery internal methods (no embedding) ---")

    # Create a mock embedding provider for get_lookup
    class MockEmbedding:
        def model_name(self): return "text-embedding-3-large"
        def embed(self, text): return result.embeddings[0]
        def embed_batch(self, texts): return result.embeddings[:len(texts)]

    query = TraceQuery(result, MockEmbedding())

    test("get_lookup()", lambda: (
        lookup := query.get_lookup(),
        assert_("axis_labels" in lookup, "should have axis_labels"),
        assert_("clusters" in lookup, "should have clusters"),
        assert_(len(lookup["clusters"]) == 10, f"expected 10 clusters, got {len(lookup['clusters'])}"),
        assert_(lookup["has_flow"] == True, "should have flow"),
        assert_(lookup["n_points"] == 348, f"expected 348 points"),
        assert_(lookup["embedding_dim"] == 3072, f"expected 3072 dims"),
        print(f"    Axis labels: {lookup['axis_labels']}"),
        print(f"    Clusters: {[c['label'] for c in lookup['clusters']]}"),
    ))

    test("get_lookup() cluster structure", lambda: (
        lookup := query.get_lookup(),
        c := lookup["clusters"][0],
        assert_("id" in c, "cluster should have id"),
        assert_("label" in c, "cluster should have label"),
        assert_("centroid_3d" in c, "cluster should have centroid_3d"),
        assert_("size" in c, "cluster should have size"),
        assert_("sample_texts" in c, "cluster should have sample_texts"),
        assert_(len(c["centroid_3d"]) == 3, "centroid should be 3D"),
    ))

    test("get_lookup() axis_ranges", lambda: (
        lookup := query.get_lookup(),
        assert_(len(lookup["axis_ranges"]) == 3, "should have 3 axis ranges"),
        r := lookup["axis_ranges"][0],
        assert_("axis" in r and "min" in r and "max" in r, "range should have axis/min/max"),
        assert_(r["min"] < r["max"], "min should be less than max"),
    ))

    test("get_lookup() flow_bounds", lambda: (
        lookup := query.get_lookup(),
        fb := lookup["flow_bounds"],
        assert_(fb["axis_min"] is not None, "should have axis_min"),
        assert_(fb["axis_max"] is not None, "should have axis_max"),
        assert_(len(fb["axis_min"]) == 3, "axis_min should be 3D"),
    ))

    # Test internal helpers with known data points
    import numpy as np
    pt = result.projected_3d[0]

    test("_compute_axis_pcts()", lambda: (
        pcts := query._compute_axis_pcts(pt),
        assert_(len(pcts) == 3, "should have 3 axis percentages"),
        assert_(all(0 <= v <= 100 for v in pcts.values()), f"percentages out of range: {pcts}"),
    ))

    test("_compute_cluster_distances()", lambda: (
        dists := query._compute_cluster_distances(pt),
        assert_(len(dists) == 10, f"should have 10 cluster distances, got {len(dists)}"),
        assert_(all(0 <= v <= 100 for v in dists.values()), f"distances out of range: {dists}"),
    ))

    test("_find_nearest_points()", lambda: (
        nearest := query._find_nearest_points(pt, k=5),
        assert_(len(nearest) == 5, f"should have 5 nearest, got {len(nearest)}"),
        assert_(nearest[0]["distance"] <= nearest[1]["distance"], "should be sorted by distance"),
    ))

    test("_sample_velocity()", lambda: (
        vel := query._sample_velocity(pt),
        assert_(vel is not None, "should have velocity"),
        assert_(len(vel) == 3, "velocity should be 3D"),
        assert_(np.isfinite(vel).all(), f"velocity should be finite: {vel}"),
    ))

    test("_decompose_vector()", lambda: (
        vel := query._sample_velocity(pt),
        decomp := query._decompose_vector(pt, vel),
        assert_("axis_decomposition" in decomp, "should have axis_decomposition"),
        assert_("cluster_pull" in decomp, "should have cluster_pull"),
        assert_("nearby_points" in decomp, "should have nearby_points"),
        assert_(len(decomp["axis_decomposition"]) == 3, "should have 3 axes"),
        assert_(len(decomp["cluster_pull"]) == 10, "should have 10 clusters"),
    ))

    # Test with mock embeddings (reuses existing embeddings, no API call)
    print("\n--- TraceQuery API methods (mock embeddings) ---")

    test("explain_path() (mock)", lambda: (
        path := query.explain_path(["text1", "text2", "text3"]),
        assert_("path_3d" in path, "should have path_3d"),
        assert_("points" in path, "should have points"),
        assert_(len(path["points"]) == 3, "should have 3 points"),
        assert_("axis_percentages" in path["points"][0], "point should have axis_percentages"),
        assert_("cluster_distances" in path["points"][0], "point should have cluster_distances"),
        assert_("nearest_texts" in path["points"][0], "point should have nearest_texts"),
    ))

    test("query_flow_at() (mock)", lambda: (
        flow := query.query_flow_at("some text"),
        assert_("position_3d" in flow, "should have position_3d"),
        assert_("velocity" in flow, "should have velocity"),
        assert_("speed" in flow, "should have speed"),
        assert_("axis_decomposition" in flow, "should have axis_decomposition"),
        assert_("cluster_pull" in flow, "should have cluster_pull"),
        assert_(flow["source"] == "flow_field", "source should be flow_field"),
    ))

    test("query_direction_at() (mock)", lambda: (
        dir_info := query.query_direction_at(["text1", "text2"]),
        assert_("estimated_direction" in dir_info, "should have estimated_direction"),
        assert_("estimated_magnitude" in dir_info, "should have estimated_magnitude"),
        assert_(dir_info["source"] == "path_estimate", "source should be path_estimate"),
    ))

    test("path_similarity() (mock)", lambda: (
        sim := query.path_similarity(["a", "b", "c"], ["d", "e", "f"]),
        assert_("frechet_distance" in sim, "should have frechet_distance"),
        assert_("mean_cosine_similarity" in sim, "should have mean_cosine_similarity"),
        assert_("direction_similarity" in sim, "should have direction_similarity"),
        assert_("start_similarity" in sim, "should have start_similarity"),
        assert_("end_similarity" in sim, "should have end_similarity"),
        assert_("overall_score" in sim, "should have overall_score"),
        assert_(0 <= sim["overall_score"] <= 1, f"overall_score out of [0,1]: {sim['overall_score']}"),
    ))

    # ── Test edge cases ──
    print("\n--- Edge cases ---")

    test("explain_path() with 1 text raises ValueError", lambda: (
        expect_error(ValueError, lambda: query.explain_path(["only one"])),
    ))

    test("query_direction_at() with 1 text raises ValueError", lambda: (
        expect_error(ValueError, lambda: query.query_direction_at(["only one"])),
    ))

    test("path_similarity() with 1-text path raises ValueError", lambda: (
        expect_error(ValueError, lambda: query.path_similarity(["one"], ["a", "b"])),
    ))

    # ── Test velocity grid sampling at boundaries ──
    print("\n--- Velocity grid boundary tests ---")

    test("velocity at exact min corner", lambda: (
        corner := np.array(result.axis_min, dtype=np.float32),
        vel := query._sample_velocity(corner),
        assert_(vel is not None and np.isfinite(vel).all(), f"corner velocity invalid: {vel}"),
    ))

    test("velocity at exact max corner", lambda: (
        corner := np.array(result.axis_max, dtype=np.float32),
        vel := query._sample_velocity(corner),
        assert_(vel is not None and np.isfinite(vel).all(), f"corner velocity invalid: {vel}"),
    ))

    test("velocity outside bounds (extrapolation)", lambda: (
        outside := np.array(result.axis_max + 10, dtype=np.float32),
        vel := query._sample_velocity(outside),
        assert_(vel is not None and np.isfinite(vel).all(), f"outside velocity invalid: {vel}"),
    ))

    # ── Test save/load roundtrip ──
    print("\n--- Save/load roundtrip ---")
    import tempfile

    test("save_result + load_result roundtrip", lambda: (
        tmp := tempfile.mktemp(prefix="ts_test_"),
        result.save_result(tmp),
        loaded := AnalysisResult.load_result(tmp),
        assert_(loaded.n_entries == result.n_entries, "entry count mismatch"),
        assert_(loaded.clusters.n_clusters == result.clusters.n_clusters, "cluster count mismatch"),
        assert_(loaded.axis_info.labels == result.axis_info.labels, "axis labels mismatch"),
        assert_(loaded.cluster_labels == result.cluster_labels, "cluster labels mismatch"),
        assert_(np.allclose(loaded.projected_3d, result.projected_3d), "projected_3d mismatch"),
        assert_(np.allclose(loaded.velocity_grid, result.velocity_grid), "velocity_grid mismatch"),
        # Cleanup
        os.remove(tmp + ".json"),
        os.remove(tmp + ".npz"),
    ))

    # ── Test importer ──
    print("\n--- Importer tests ---")

    from tracescope.analysis.importer import from_list, from_lists, auto_import

    test("from_list() basic", lambda: (
        session := from_list(["hello", "world", "test"]),
        assert_(len(session) == 3, f"expected 3 entries, got {len(session)}"),
        assert_(session.entries[0].text == "hello", "text mismatch"),
        assert_(session.entries[0].role == "entry", f"role should be 'entry', got '{session.entries[0].role}'"),
    ))

    test("from_lists() multi-path", lambda: (
        session := from_lists([["path1a", "path1b"], ["path2a", "path2b", "path2c"]]),
        assert_(len(session) == 5, f"expected 5 entries, got {len(session)}"),
        assert_(session.entries[0].path_id == 0, "first entry path_id should be 0"),
        assert_(session.entries[2].path_id == 1, "third entry path_id should be 1"),
    ))

    test("from_list() empty raises", lambda: (
        expect_error(ValueError, lambda: from_list([])),
    ))

    test("auto_import() news_headlines.json", lambda: (
        session := auto_import(os.path.join(os.path.dirname(__file__), "..", "examples", "sample_data", "news_headlines.json")),
        assert_(len(session) > 0, "should have entries"),
    ))

    # ── Test config ──
    print("\n--- Config tests ---")

    from tracescope.config import TraceScopeConfig

    test("TraceScopeConfig defaults", lambda: (
        cfg := TraceScopeConfig(),
        assert_(cfg.embedding_model == "text-embedding-3-large", f"default model: {cfg.embedding_model}"),
        assert_(cfg.chroma_path is not None, "should have chroma_path"),
    ))

    test("TraceScopeConfig with custom values", lambda: (
        cfg := TraceScopeConfig(
            openai_api_key="sk-test",
            embedding_model="text-embedding-3-small",
            llm_model="gpt-4o",
        ),
        assert_(cfg.openai_api_key == "sk-test"),
        assert_(cfg.embedding_model == "text-embedding-3-small"),
        assert_(cfg.llm_model == "gpt-4o"),
    ))

    # ── Test flow field system (no API calls needed) ──
    print("\n--- FlowFieldSystem tests ---")

    from tracescope.visualization.flow_field import FlowFieldSystem

    test("FlowFieldSystem init with velocity grid", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=10,
        ),
        assert_(flow.positions.shape[1] == 3, "positions should be 3D"),
        assert_(len(flow.positions) > 0, "should have particles"),
    ))

    test("FlowFieldSystem step()", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        flow.step(),
        assert_(flow.positions.shape[1] == 3, "positions should be 3D after step"),
    ))

    test("FlowFieldSystem sample_velocity_batch()", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        pts := result.projected_3d[:10],
        vels := flow.sample_velocity_batch(pts),
        assert_(vels.shape == (10, 3), f"expected (10,3), got {vels.shape}"),
        assert_(np.isfinite(vels).all(), "velocities should be finite"),
    ))

    test("FlowFieldSystem with path_points (blob)", lambda: (
        path_pts := result.projected_3d[:20].astype(np.float32),
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=10,
            path_points=path_pts,
        ),
        assert_(flow._path_mask is not None, "should have path mask"),
        assert_(flow._path_mask.shape == (flow.BLOB_RESOLUTION,) * 3, "mask shape mismatch"),
    ))

    test("FlowFieldSystem ball tracking", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        mid := (result.axis_min + result.axis_max) / 2,
        flow.set_ball_position(mid[0], mid[1], mid[2]),
        flow.advance_ball(),
        assert_(len(flow.ball_trail) > 0, "ball should have trail"),
    ))

    test("FlowFieldSystem precompute_frames()", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        frames := flow.precompute_frames(3),
        assert_(len(frames) == 3, f"expected 3 frames, got {len(frames)}"),
        assert_("positions" in frames[0], "frame should have positions"),
        assert_("colors" in frames[0], "frame should have colors"),
    ))

    # ── Summary ──
    print("\n" + "=" * 60)
    print(f"RESULTS: {passed} passed, {failed} failed")
    if errors:
        print("\nFailed tests:")
        for name, err in errors:
            print(f"  - {name}: {err}")
    print("=" * 60)

    return 0 if failed == 0 else 1


def assert_(condition, msg="assertion failed"):
    if not condition:
        raise AssertionError(msg)


def expect_error(exc_type, fn):
    try:
        fn()
        raise AssertionError(f"Expected {exc_type.__name__} but no error was raised")
    except exc_type:
        pass


if __name__ == "__main__":
    sys.exit(main())
