"""
Stress test edge cases — no API calls, uses cached data and synthetic inputs.

Run: python test/test_edge_cases.py
"""

import sys
import os
import json
import tempfile
import traceback

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

import numpy as np
from tracescope.models.analysis import AnalysisResult, ClusterResult, AxisInfo
from tracescope.models.trace import TraceEntry, TraceSession
from tracescope.analysis.importer import from_list, from_lists, auto_import
from tracescope.visualization.flow_field import FlowFieldSystem
from tracescope.query import TraceQuery

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 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


def main():
    global passed, failed

    print("=" * 60)
    print("TraceScope Edge Case Stress Tests")
    print("=" * 60)

    result = AnalysisResult.load_result(CACHE_PATH)

    # ── Importer edge cases ──
    print("\n--- Importer edge cases ---")

    test("from_list() single text", lambda: (
        s := from_list(["only one"]),
        assert_(len(s) == 1, f"expected 1, got {len(s)}"),
    ))

    test("from_list() with empty strings", lambda: (
        s := from_list(["", "hello", ""]),
        assert_(len(s) == 3, "should keep empty strings"),
    ))

    test("from_list() very long text", lambda: (
        s := from_list(["x" * 100000]),
        assert_(len(s) == 1, "should handle very long text"),
        assert_(len(s.entries[0].text) == 100000, "text should be preserved"),
    ))

    test("from_list() unicode text", lambda: (
        s := from_list(["こんにちは世界", "مرحبا بالعالم", "🚀🌍🎉", "Ñoño"]),
        assert_(len(s) == 4, "should handle unicode"),
    ))

    test("from_lists() single path", lambda: (
        s := from_lists([["a", "b", "c"]]),
        assert_(len(s) == 3, f"expected 3, got {len(s)}"),
        assert_(all(e.path_id == 0 for e in s.entries), "all should be path 0"),
    ))

    test("from_lists() many paths", lambda: (
        paths := [[f"path{i}_msg{j}" for j in range(3)] for i in range(100)],
        s := from_lists(paths),
        assert_(len(s) == 300, f"expected 300, got {len(s)}"),
    ))

    test("from_lists() empty path in list", lambda: (
        expect_error((ValueError, IndexError), lambda: from_lists([["a", "b"], [], ["c"]])),
    ))

    test("auto_import() nonexistent file", lambda: (
        expect_error((FileNotFoundError, Exception), lambda: auto_import("/nonexistent/file.json")),
    ))

    # Test auto_import with different formats
    test("auto_import() OpenAI chat format", lambda: (
        tmp := tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False),
        json.dump({"messages": [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there"},
        ]}, tmp),
        tmp.close(),
        s := auto_import(tmp.name),
        assert_(len(s) >= 2, f"expected >=2, got {len(s)}"),
        os.unlink(tmp.name),
    ))

    test("auto_import() plain text list", lambda: (
        tmp := tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False),
        json.dump(["text one", "text two", "text three"], tmp),
        tmp.close(),
        s := auto_import(tmp.name),
        assert_(len(s) == 3, f"expected 3, got {len(s)}"),
        os.unlink(tmp.name),
    ))

    test("auto_import() Anthropic format", lambda: (
        tmp := tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False),
        json.dump([
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": [{"type": "text", "text": "Hi there"}]},
        ], tmp),
        tmp.close(),
        s := auto_import(tmp.name),
        assert_(len(s) >= 2, f"expected >=2, got {len(s)}"),
        os.unlink(tmp.name),
    ))

    # ── FlowFieldSystem edge cases ──
    print("\n--- FlowFieldSystem edge cases ---")

    test("FlowFieldSystem with particle_grid=1", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=1,
        ),
        assert_(len(flow.positions) == 1, f"expected 1 particle, got {len(flow.positions)}"),
    ))

    test("FlowFieldSystem with particle_grid=50 (high)", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=50,
        ),
        assert_(len(flow.positions) == 50**3, f"expected {50**3}, got {len(flow.positions)}"),
    ))

    test("FlowFieldSystem 100 steps stability", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        [flow.step() for _ in range(100)],
        assert_(np.isfinite(flow.positions).all(), "positions should be finite after 100 steps"),
    ))

    test("FlowFieldSystem ball at boundary", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        flow.set_ball_position(result.axis_min[0], result.axis_min[1], result.axis_min[2]),
        [flow.advance_ball() for _ in range(50)],
        assert_(np.isfinite(flow.ball_pos).all(), "ball pos should be finite"),
    ))

    test("FlowFieldSystem sample_velocity at grid corners", lambda: (
        flow := FlowFieldSystem(
            result.velocity_grid, result.axis_min, result.axis_max,
            particle_grid=5,
        ),
        v1 := flow.sample_velocity(result.axis_min[0], result.axis_min[1], result.axis_min[2]),
        v2 := flow.sample_velocity(result.axis_max[0], result.axis_max[1], result.axis_max[2]),
        assert_(np.isfinite(v1).all(), f"min corner velocity: {v1}"),
        assert_(np.isfinite(v2).all(), f"max corner velocity: {v2}"),
    ))

    test("FlowFieldSystem with zero velocity grid", lambda: (
        zero_grid := np.zeros((40, 40, 40, 3), dtype=np.float32),
        flow := FlowFieldSystem(zero_grid, result.axis_min, result.axis_max, particle_grid=5),
        flow.step(),
        # Particles shouldn't move with zero velocity
        # (they still wrap/respawn due to aging)
    ))

    # ── Velocity grid data quality ──
    print("\n--- Velocity grid data quality ---")

    test("velocity grid finite everywhere", lambda: (
        assert_(np.isfinite(result.velocity_grid).all(), "velocity grid has NaN/Inf"),
    ))

    test("velocity grid not all zeros", lambda: (
        assert_(np.abs(result.velocity_grid).sum() > 0, "velocity grid is all zeros"),
    ))

    test("velocity grid magnitude distribution", lambda: (
        mags := np.linalg.norm(result.velocity_grid, axis=-1),
        print(f"    Min speed: {mags.min():.6f}"),
        print(f"    Max speed: {mags.max():.6f}"),
        print(f"    Mean speed: {mags.mean():.6f}"),
        print(f"    Std speed: {mags.std():.6f}"),
        print(f"    Zero-speed cells: {(mags == 0).sum()} / {mags.size}"),
    ))

    # ── Cluster consistency ──
    print("\n--- Cluster consistency ---")

    test("all entries have valid cluster labels", lambda: (
        labels := result.clusters.labels,
        assert_(all(0 <= l < result.clusters.n_clusters for l in labels),
                f"invalid cluster labels: min={min(labels)}, max={max(labels)}"),
    ))

    def check_centroids():
        for c in range(result.clusters.n_clusters):
            indices = [i for i, l in enumerate(result.clusters.labels) if l == c]
            if indices:
                actual_center = result.projected_3d[indices].mean(axis=0)
                stored_center = result.cluster_centroids_3d[c]
                dist = np.linalg.norm(actual_center - stored_center)
                if dist > 1.0:
                    print(f"    WARNING: Cluster {c} centroid drift: {dist:.4f}")
    test("cluster centroids match actual cluster centers", check_centroids)

    test("cluster labels count matches n_clusters", lambda: (
        assert_(len(result.cluster_labels) == result.clusters.n_clusters,
                f"labels: {len(result.cluster_labels)}, clusters: {result.clusters.n_clusters}"),
    ))

    # ── Projection consistency ──
    print("\n--- Projection consistency ---")

    def check_bounds():
        pts = result.projected_3d
        actual_min = pts.min(axis=0)
        actual_max = pts.max(axis=0)
        for i in range(3):
            assert_(result.axis_min[i] <= actual_min[i] + 0.1,
                    f"axis_min[{i}]={result.axis_min[i]} > data min={actual_min[i]}")
            assert_(result.axis_max[i] >= actual_max[i] - 0.1,
                    f"axis_max[{i}]={result.axis_max[i]} < data max={actual_max[i]}")
    test("projected_3d matches axis_min/max bounds", check_bounds)

    test("no duplicate 3D projections", lambda: (
        pts := result.projected_3d,
        n_unique := len(np.unique(pts, axis=0)),
        dup_pct := (1 - n_unique / len(pts)) * 100,
        print(f"    Unique points: {n_unique}/{len(pts)} ({dup_pct:.1f}% duplicates)"),
        # Some duplicates are OK (e.g., very similar texts)
        assert_(n_unique > len(pts) * 0.5, f"too many duplicate projections: {dup_pct:.1f}%"),
    ))

    # ── Save/load edge cases ──
    print("\n--- Save/load edge cases ---")

    test("save/load preserves all arrays exactly", lambda: (
        tmp := tempfile.mktemp(prefix="ts_roundtrip_"),
        result.save_result(tmp),
        loaded := AnalysisResult.load_result(tmp),
        assert_(np.array_equal(loaded.embeddings, result.embeddings), "embeddings mismatch"),
        assert_(np.array_equal(loaded.projected_3d, result.projected_3d), "projected_3d mismatch"),
        assert_(np.array_equal(loaded.velocity_grid, result.velocity_grid), "velocity_grid mismatch"),
        assert_(np.array_equal(loaded.axis_min, result.axis_min), "axis_min mismatch"),
        assert_(np.array_equal(loaded.axis_max, result.axis_max), "axis_max mismatch"),
        os.remove(tmp + ".json"),
        os.remove(tmp + ".npz"),
    ))

    test("load_result with missing npz file", lambda: (
        expect_error((FileNotFoundError, Exception), lambda: AnalysisResult.load_result("/nonexistent/path")),
    ))

    # ── DTW / Frechet edge cases ──
    print("\n--- DTW / Frechet edge cases ---")

    class MockEmb:
        def model_name(self): return "test"
        def embed(self, text): return result.embeddings[0]
        def embed_batch(self, texts): return result.embeddings[:len(texts)]

    query = TraceQuery(result, MockEmb())

    test("path_similarity identical paths", lambda: (
        sim := query.path_similarity(["a", "b"], ["a", "b"]),
        # With mock embeddings, both paths use same vectors, so similarity should be high
        print(f"    Identical path score: {sim['overall_score']:.4f}"),
        print(f"    Frechet distance: {sim['frechet_distance']:.6f}"),
    ))

    test("path_similarity 2-text paths", lambda: (
        sim := query.path_similarity(["a", "b"], ["c", "d"]),
        assert_(0 <= sim["overall_score"] <= 1, f"score out of range: {sim['overall_score']}"),
    ))

    test("path_similarity asymmetric lengths", lambda: (
        sim := query.path_similarity(["a", "b"], ["c", "d", "e", "f", "g"]),
        assert_(0 <= sim["overall_score"] <= 1, f"score out of range: {sim['overall_score']}"),
    ))

    # ── Hashing edge cases ──
    print("\n--- Hashing edge cases ---")

    from tracescope.utils.hashing import sha256_key

    test("sha256_key deterministic", lambda: (
        k1 := sha256_key("model", "system", "user"),
        k2 := sha256_key("model", "system", "user"),
        assert_(k1 == k2, "same inputs should give same hash"),
    ))

    test("sha256_key different for different inputs", lambda: (
        k1 := sha256_key("model1", "sys", "user"),
        k2 := sha256_key("model2", "sys", "user"),
        assert_(k1 != k2, "different inputs should give different hash"),
    ))

    # ── Cache edge cases ──
    print("\n--- Cache edge cases ---")

    from tracescope.storage.cache import LLMResponseCache, ResultCache

    test("LLMResponseCache disabled mode", lambda: (
        cache := LLMResponseCache("", enabled=False),
        assert_(cache.get("m", "s", "u") is None, "disabled cache should return None"),
        cache.put("m", "s", "u", "response"),  # should be no-op
        assert_(cache.get("m", "s", "u") is None, "disabled cache should still return None"),
    ))

    test("LLMResponseCache in-memory roundtrip", lambda: (
        tmp := tempfile.NamedTemporaryFile(suffix=".db", delete=False),
        tmp.close(),
        cache := LLMResponseCache(tmp.name),
        cache.put("model", "system", "user", "response_text"),
        got := cache.get("model", "system", "user"),
        assert_(got == "response_text", f"expected 'response_text', got '{got}'"),
        # Different prompt should miss
        assert_(cache.get("model", "system", "different") is None, "should miss on different prompt"),
        cache.close(),
        os.unlink(tmp.name),
    ))

    # ── 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


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