"""
Enron Emails Benchmark Evaluation - Retrieval Recall@3 (Stages 2/3/4)
Follows the style of FinanceBench/LAION: Stage 2 recall vs FAISS baseline,
Stage 3 complexity sweep to target recall, Stage 4 index comparison.
On errors, fail fast without fallbacks.
"""

import argparse
import json
import os
import pickle

import numpy as np
from leann import LeannBuilder, LeannSearcher
from leann_backend_hnsw import faiss


class RecallEvaluator:
    """Stage 2: Evaluate Recall@3 (LEANN vs FAISS)"""

    def __init__(self, index_path: str, baseline_dir: str):
        self.index_path = index_path
        self.baseline_dir = baseline_dir
        self.searcher = LeannSearcher(index_path)

        baseline_index_path = os.path.join(baseline_dir, "faiss_flat.index")
        metadata_path = os.path.join(baseline_dir, "metadata.pkl")

        self.faiss_index = faiss.read_index(baseline_index_path)
        with open(metadata_path, "rb") as f:
            self.passage_ids = pickle.load(f)

        print(f"📚 Loaded FAISS flat baseline with {self.faiss_index.ntotal} vectors")

        # No fallbacks here; if embedding server is needed but fails, the caller will see the error.

    def evaluate_recall_at_3(
        self, queries: list[str], complexity: int = 64, recompute_embeddings: bool = True
    ) -> float:
        """Evaluate recall@3 using FAISS Flat as ground truth"""
        from leann.api import compute_embeddings

        recompute_str = "with recompute" if recompute_embeddings else "no recompute"
        print(f"🔍 Evaluating recall@3 with complexity={complexity} ({recompute_str})...")

        total_recall = 0.0
        for i, query in enumerate(queries):
            # Compute query embedding with the same model/mode as the index
            q_emb = compute_embeddings(
                [query],
                self.searcher.embedding_model,
                mode=self.searcher.embedding_mode,
                use_server=False,
            ).astype(np.float32)

            # Search FAISS Flat ground truth
            n = q_emb.shape[0]
            k = 3
            distances = np.zeros((n, k), dtype=np.float32)
            labels = np.zeros((n, k), dtype=np.int64)
            self.faiss_index.search(
                n,
                faiss.swig_ptr(q_emb),
                k,
                faiss.swig_ptr(distances),
                faiss.swig_ptr(labels),
            )

            baseline_ids = {self.passage_ids[idx] for idx in labels[0]}

            # Search with LEANN (may require embedding server depending on index configuration)
            results = self.searcher.search(
                query,
                top_k=3,
                complexity=complexity,
                recompute_embeddings=recompute_embeddings,
            )
            test_ids = {r.id for r in results}

            intersection = test_ids.intersection(baseline_ids)
            recall = len(intersection) / 3.0
            total_recall += recall

            if i < 3:
                print(f"  Q{i + 1}: '{query[:60]}...' -> Recall@3: {recall:.3f}")
                print(f"    FAISS: {list(baseline_ids)}")
                print(f"    LEANN: {list(test_ids)}")
                print(f"    ∩: {list(intersection)}")

        avg = total_recall / max(1, len(queries))
        print(f"📊 Average Recall@3: {avg:.3f} ({avg * 100:.1f}%)")
        return avg

    def cleanup(self):
        if hasattr(self, "searcher"):
            self.searcher.cleanup()


class EnronEvaluator:
    def __init__(self, index_path: str):
        self.index_path = index_path
        self.searcher = LeannSearcher(index_path)

    def load_queries(self, queries_file: str) -> list[str]:
        queries: list[str] = []
        with open(queries_file, encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                data = json.loads(line)
                if "query" in data:
                    queries.append(data["query"])
        print(f"📊 Loaded {len(queries)} queries from {queries_file}")
        return queries

    def cleanup(self):
        if self.searcher:
            self.searcher.cleanup()

    def analyze_index_sizes(self) -> dict:
        """Analyze index sizes (.index only), similar to LAION bench."""
        from pathlib import Path

        print("📏 Analyzing index sizes (.index only)...")
        index_path = Path(self.index_path)
        index_dir = index_path.parent
        index_name = index_path.stem

        sizes: dict[str, float] = {}
        index_file = index_dir / f"{index_name}.index"
        meta_file = index_dir / f"{index_path.name}.meta.json"
        passages_file = index_dir / f"{index_path.name}.passages.jsonl"
        passages_idx_file = index_dir / f"{index_path.name}.passages.idx"

        sizes["index_only_mb"] = (
            index_file.stat().st_size / (1024 * 1024) if index_file.exists() else 0.0
        )
        sizes["metadata_mb"] = (
            meta_file.stat().st_size / (1024 * 1024) if meta_file.exists() else 0.0
        )
        sizes["passages_text_mb"] = (
            passages_file.stat().st_size / (1024 * 1024) if passages_file.exists() else 0.0
        )
        sizes["passages_index_mb"] = (
            passages_idx_file.stat().st_size / (1024 * 1024) if passages_idx_file.exists() else 0.0
        )

        print(f"  📁 .index size: {sizes['index_only_mb']:.1f} MB")
        return sizes

    def create_non_compact_index_for_comparison(self, non_compact_index_path: str) -> dict:
        """Create a non-compact index for comparison using current passages and embeddings."""
        from pathlib import Path

        current_index_path = Path(self.index_path)
        current_index_dir = current_index_path.parent
        current_index_name = current_index_path.name

        # Read metadata to get passage source and embedding model
        meta_path = current_index_dir / f"{current_index_name}.meta.json"
        with open(meta_path, encoding="utf-8") as f:
            meta = json.load(f)

        passage_source = meta["passage_sources"][0]
        passage_file = passage_source["path"]

        # Convert relative path to absolute
        if not Path(passage_file).is_absolute():
            passage_file = current_index_dir / Path(passage_file).name

        # Load all passages and ids
        ids: list[str] = []
        texts: list[str] = []
        with open(passage_file, encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    data = json.loads(line)
                    ids.append(str(data["id"]))
                    texts.append(data["text"])

        # Compute embeddings using the same method as LEANN
        from leann.api import compute_embeddings

        embeddings = compute_embeddings(
            texts,
            meta["embedding_model"],
            mode=meta.get("embedding_mode", "sentence-transformers"),
            use_server=False,
        ).astype(np.float32)

        # Build non-compact index with same passages and embeddings
        builder = LeannBuilder(
            backend_name="hnsw",
            embedding_model=meta["embedding_model"],
            embedding_mode=meta.get("embedding_mode", "sentence-transformers"),
            is_recompute=False,
            is_compact=False,
            **{
                k: v
                for k, v in meta.get("backend_kwargs", {}).items()
                if k not in ["is_recompute", "is_compact"]
            },
        )

        # Persist a pickle for build_index_from_embeddings
        pkl_path = current_index_dir / f"{Path(non_compact_index_path).stem}_embeddings.pkl"
        with open(pkl_path, "wb") as pf:
            pickle.dump((ids, embeddings), pf)

        print(
            f"🔨 Building non-compact index at {non_compact_index_path} from precomputed embeddings..."
        )
        builder.build_index_from_embeddings(non_compact_index_path, str(pkl_path))

        # Analyze the non-compact index size
        temp_evaluator = EnronEvaluator(non_compact_index_path)
        non_compact_sizes = temp_evaluator.analyze_index_sizes()
        non_compact_sizes["index_type"] = "non_compact"

        return non_compact_sizes

    def compare_index_performance(
        self, non_compact_path: str, compact_path: str, test_queries: list[str], complexity: int
    ) -> dict:
        """Compare search speed for non-compact vs compact indexes."""
        import time

        results: dict = {
            "non_compact": {"search_times": []},
            "compact": {"search_times": []},
            "avg_search_times": {},
            "speed_ratio": 0.0,
        }

        print("⚡ Comparing search performance between indexes...")
        # Non-compact (no recompute)
        print("  🔍 Testing non-compact index (no recompute)...")
        non_compact_searcher = LeannSearcher(non_compact_path)
        for q in test_queries:
            t0 = time.time()
            _ = non_compact_searcher.search(
                q, top_k=3, complexity=complexity, recompute_embeddings=False
            )
            results["non_compact"]["search_times"].append(time.time() - t0)

        # Compact (with recompute). Fail fast if it cannot run.
        print("  🔍 Testing compact index (with recompute)...")
        compact_searcher = LeannSearcher(compact_path)
        for q in test_queries:
            t0 = time.time()
            _ = compact_searcher.search(
                q, top_k=3, complexity=complexity, recompute_embeddings=True
            )
            results["compact"]["search_times"].append(time.time() - t0)
        compact_searcher.cleanup()

        if results["non_compact"]["search_times"]:
            results["avg_search_times"]["non_compact"] = sum(
                results["non_compact"]["search_times"]
            ) / len(results["non_compact"]["search_times"])
        if results["compact"]["search_times"]:
            results["avg_search_times"]["compact"] = sum(results["compact"]["search_times"]) / len(
                results["compact"]["search_times"]
            )
        if results["avg_search_times"].get("compact", 0) > 0:
            results["speed_ratio"] = (
                results["avg_search_times"]["non_compact"] / results["avg_search_times"]["compact"]
            )
        else:
            results["speed_ratio"] = 0.0

        non_compact_searcher.cleanup()
        return results

    def evaluate_complexity(
        self,
        recall_eval: "RecallEvaluator",
        queries: list[str],
        target: float = 0.90,
        c_min: int = 8,
        c_max: int = 256,
        max_iters: int = 10,
        recompute: bool = False,
    ) -> dict:
        """Binary search minimal complexity achieving target recall (monotonic assumption)."""

        def round_c(x: int) -> int:
            # snap to multiple of 8 like other benches typically do
            return max(1, int((x + 7) // 8) * 8)

        metrics: list[dict] = []

        lo = round_c(c_min)
        hi = round_c(c_max)

        print(
            f"🧪 Binary search complexity in [{lo}, {hi}] for target Recall@3>={int(target * 100)}%..."
        )

        # Ensure upper bound can reach target; expand if needed (up to a cap)
        r_lo = recall_eval.evaluate_recall_at_3(
            queries, complexity=lo, recompute_embeddings=recompute
        )
        metrics.append({"complexity": lo, "recall_at_3": r_lo})
        r_hi = recall_eval.evaluate_recall_at_3(
            queries, complexity=hi, recompute_embeddings=recompute
        )
        metrics.append({"complexity": hi, "recall_at_3": r_hi})

        cap = 1024
        while r_hi < target and hi < cap:
            lo = hi
            r_lo = r_hi
            hi = round_c(hi * 2)
            r_hi = recall_eval.evaluate_recall_at_3(
                queries, complexity=hi, recompute_embeddings=recompute
            )
            metrics.append({"complexity": hi, "recall_at_3": r_hi})

        if r_hi < target:
            print(f"⚠️ Max complexity {hi} did not reach target recall {target:.2f}.")
            print("📈 Observations:")
            for m in metrics:
                print(f"  C={m['complexity']:>4} -> Recall@3={m['recall_at_3'] * 100:.1f}%")
            return {"metrics": metrics, "best_complexity": None, "target_recall": target}

        # Binary search within [lo, hi]
        best = hi
        iters = 0
        while lo < hi and iters < max_iters:
            mid = round_c((lo + hi) // 2)
            r_mid = recall_eval.evaluate_recall_at_3(
                queries, complexity=mid, recompute_embeddings=recompute
            )
            metrics.append({"complexity": mid, "recall_at_3": r_mid})
            if r_mid >= target:
                best = mid
                hi = mid
            else:
                lo = mid + 8  # move past mid, respecting multiple-of-8 step
            iters += 1

        print("📈 Binary search results (sampled points):")
        # Print unique complexity entries ordered by complexity
        for m in sorted(
            {m["complexity"]: m for m in metrics}.values(), key=lambda x: x["complexity"]
        ):
            print(f"  C={m['complexity']:>4} -> Recall@3={m['recall_at_3'] * 100:.1f}%")
        print(f"✅ Minimal complexity achieving {int(target * 100)}% recall: {best}")
        return {"metrics": metrics, "best_complexity": best, "target_recall": target}


def main():
    parser = argparse.ArgumentParser(description="Enron Emails Benchmark Evaluation")
    parser.add_argument("--index", required=True, help="Path to LEANN index")
    parser.add_argument(
        "--queries", default="data/evaluation_queries.jsonl", help="Path to evaluation queries"
    )
    parser.add_argument(
        "--stage",
        choices=["2", "3", "4", "all"],
        default="all",
        help="Which stage to run (2=recall, 3=complexity, 4=index comparison)",
    )
    parser.add_argument("--complexity", type=int, default=None, help="LEANN search complexity")
    parser.add_argument("--baseline-dir", default="baseline", help="Baseline output directory")
    parser.add_argument(
        "--max-queries", type=int, help="Limit number of queries to evaluate", default=1000
    )
    parser.add_argument(
        "--target-recall", type=float, default=0.90, help="Target Recall@3 for Stage 3"
    )
    parser.add_argument("--output", help="Save results to JSON file")

    args = parser.parse_args()

    # Resolve queries file: if default path not found, fall back to index's directory
    if not os.path.exists(args.queries):
        from pathlib import Path

        idx_dir = Path(args.index).parent
        fallback_q = idx_dir / "evaluation_queries.jsonl"
        if fallback_q.exists():
            args.queries = str(fallback_q)

    baseline_index_path = os.path.join(args.baseline_dir, "faiss_flat.index")
    if not os.path.exists(baseline_index_path):
        print(f"❌ FAISS baseline not found at {baseline_index_path}")
        print("💡 Please run setup_enron_emails.py first to build the baseline")
        raise SystemExit(1)

    results_out: dict = {}

    if args.stage in ("2", "all"):
        print("🚀 Starting Stage 2: Recall@3 evaluation")
        evaluator = RecallEvaluator(args.index, args.baseline_dir)

        enron_eval = EnronEvaluator(args.index)
        queries = enron_eval.load_queries(args.queries)
        queries = queries[:10]
        print(f"🧪 Using first {len(queries)} queries")

        complexity = args.complexity or 64
        r = evaluator.evaluate_recall_at_3(queries, complexity)
        results_out["stage2"] = {"complexity": complexity, "recall_at_3": r}
        evaluator.cleanup()
        enron_eval.cleanup()
        print("✅ Stage 2 completed!\n")

    if args.stage in ("3", "all"):
        print("🚀 Starting Stage 3: Binary search for target recall (no recompute)")
        enron_eval = EnronEvaluator(args.index)
        queries = enron_eval.load_queries(args.queries)
        queries = queries[: args.max_queries]
        print(f"🧪 Using first {len(queries)} queries")

        # Build non-compact index for fast binary search (recompute_embeddings=False)
        from pathlib import Path
        index_path = Path(args.index)
        non_compact_index_path = str(index_path.parent / f"{index_path.stem}_noncompact.leann")
        enron_eval.create_non_compact_index_for_comparison(non_compact_index_path)

        # Use non-compact evaluator for binary search with recompute=False
        evaluator_nc = RecallEvaluator(non_compact_index_path, args.baseline_dir)
        sweep = enron_eval.evaluate_complexity(
            evaluator_nc, queries, target=args.target_recall, recompute=False
        )
        results_out["stage3"] = sweep
        # Persist default stage 3 results near the index for Stage 4 auto-pickup
        from pathlib import Path
        default_stage3_path = Path(args.index).parent / "enron_stage3_results.json"
        with open(default_stage3_path, "w", encoding="utf-8") as f:
            json.dump({"stage3": sweep}, f, indent=2)
        print(f"📝 Saved Stage 3 summary to {default_stage3_path}")
        evaluator_nc.cleanup()
        enron_eval.cleanup()
        print("✅ Stage 3 completed!\n")

    if args.stage in ("4", "all"):
        print("🚀 Starting Stage 4: Index size + performance comparison")
        evaluator = RecallEvaluator(args.index, args.baseline_dir)
        enron_eval = EnronEvaluator(args.index)
        queries = enron_eval.load_queries(args.queries)
        test_q = queries[: min(args.max_queries, len(queries))]

        current_sizes = enron_eval.analyze_index_sizes()
        # Build non-compact index for comparison (no fallback)
        from pathlib import Path

        index_path = Path(args.index)
        non_compact_path = str(index_path.parent / f"{index_path.stem}_noncompact.leann")
        non_compact_sizes = enron_eval.create_non_compact_index_for_comparison(non_compact_path)
        nc_eval = EnronEvaluator(non_compact_path)

        if (
            current_sizes.get("index_only_mb", 0) > 0
            and non_compact_sizes.get("index_only_mb", 0) > 0
        ):
            storage_saving_percent = max(
                0.0,
                100.0 * (1.0 - current_sizes["index_only_mb"] / non_compact_sizes["index_only_mb"]),
            )
        else:
            storage_saving_percent = 0.0

        if args.complexity is None:
            # Prefer in-session Stage 3 result
            if "stage3" in results_out and results_out["stage3"].get("best_complexity") is not None:
                complexity = results_out["stage3"]["best_complexity"]
                print(f"📥 Using best complexity from Stage 3 in-session: {complexity}")
            else:
                # Try to load last saved Stage 3 result near index
                default_stage3_path = Path(args.index).parent / "enron_stage3_results.json"
                if default_stage3_path.exists():
                    with open(default_stage3_path, encoding="utf-8") as f:
                        prev = json.load(f)
                    complexity = prev.get("stage3", {}).get("best_complexity")
                    if complexity is None:
                        raise SystemExit("❌ Stage 4: No --complexity and no best_complexity found in saved Stage 3 results")
                    print(f"📥 Using best complexity from saved Stage 3: {complexity}")
                else:
                    raise SystemExit(
                        "❌ Stage 4 requires --complexity if Stage 3 hasn't been run. Run stage 3 first or pass --complexity."
                    )
        else:
            complexity = args.complexity

        comp = enron_eval.compare_index_performance(
            non_compact_path, args.index, test_q, complexity=complexity
        )
        results_out["stage4"] = {
            "current_index": current_sizes,
            "non_compact_index": non_compact_sizes,
            "storage_saving_percent": storage_saving_percent,
            "performance_comparison": comp,
        }
        nc_eval.cleanup()
        evaluator.cleanup()
        enron_eval.cleanup()
        print("✅ Stage 4 completed!\n")

    if args.output and results_out:
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(results_out, f, indent=2)
        print(f"📝 Saved results to {args.output}")


if __name__ == "__main__":
    main()