#!/usr/bin/env python3
  """
  Compare two Elasticsearch indices:
  - mapping structure (field paths + types)
  - field coverage stats (exists; nested-safe)
  - random sample documents (same _id) and diff _source field paths
  
  Usage:
    python scripts/inspect/compare_indices.py INDEX_A INDEX_B --sample-size 25
    python scripts/inspect/compare_indices.py INDEX_A INDEX_B --fields title.zh,vendor.zh,keywords.zh,tags.zh --fields-nested image_embedding.url,enriched_attributes.name
  """
  
  from __future__ import annotations
  
  import argparse
  import json
  import sys
  from dataclasses import dataclass
  from pathlib import Path
  from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
  
  sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
  
  from utils.es_client import ESClient, get_es_client_from_env
  
  
  def _walk_mapping_properties(props: Dict[str, Any], prefix: str = "") -> Dict[str, str]:
      """Flatten mapping properties into {field_path: type} including multi-fields."""
      out: Dict[str, str] = {}
      for name, node in (props or {}).items():
          path = f"{prefix}.{name}" if prefix else name
          if not isinstance(node, dict):
              out[path] = "unknown"
              continue
          out[path] = node.get("type") or "object"
          if isinstance(node.get("properties"), dict):
              out.update(_walk_mapping_properties(node["properties"], path))
          if isinstance(node.get("fields"), dict):
              for sub, subnode in node["fields"].items():
                  if isinstance(subnode, dict):
                      out[f"{path}.{sub}"] = subnode.get("type") or "object"
                  else:
                      out[f"{path}.{sub}"] = "unknown"
      return out
  
  
  def _get_top_level_field_type(mapping: Dict[str, Any], top_field: str) -> Optional[str]:
      props = mapping.get("mappings", {}).get("properties", {}) or {}
      node = props.get(top_field)
      if not isinstance(node, dict):
          return None
      return node.get("type") or "object"
  
  
  def _field_paths_from_source(obj: Any, prefix: str = "", list_depth: int = 3) -> Set[str]:
      """Return dotted field paths found in _source. For lists, uses '[]' marker."""
      out: Set[str] = set()
      if isinstance(obj, dict):
          for k, v in obj.items():
              p = f"{prefix}.{k}" if prefix else k
              out.add(p)
              out |= _field_paths_from_source(v, p, list_depth=list_depth)
      elif isinstance(obj, list):
          # Do not explode: just traverse first N elements
          for v in obj[:list_depth]:
              p = f"{prefix}[]" if prefix else "[]"
              out |= _field_paths_from_source(v, p, list_depth=list_depth)
      return out
  
  
  def _chunks(seq: List[str], size: int) -> Iterable[List[str]]:
      for i in range(0, len(seq), size):
          yield seq[i : i + size]
  
  
  @dataclass(frozen=True)
  class CoverageField:
      field: str
      # If set, use nested query with this path (e.g. "image_embedding").
      nested_path: Optional[str] = None
  
  
  def _infer_coverage_fields(
      mapping: Dict[str, Any],
      raw_fields: List[str],
      raw_nested_fields: List[str],
  ) -> List[CoverageField]:
      """
      Build coverage fields list. For fields in raw_nested_fields, always treat as nested
      and infer nested path as first segment.
      For raw_fields, auto-detect nested by checking mapping top-level field type.
      """
      out: List[CoverageField] = []
  
      nested_set = {f.strip() for f in raw_nested_fields if f.strip()}
      for f in nested_set:
          path = f.split(".", 1)[0]
          out.append(CoverageField(field=f, nested_path=path))
  
      for f in [x.strip() for x in raw_fields if x.strip()]:
          if f in nested_set:
              continue
          top = f.split(".", 1)[0]
          top_type = _get_top_level_field_type(mapping, top)
          if top_type == "nested":
              out.append(CoverageField(field=f, nested_path=top))
          else:
              out.append(CoverageField(field=f, nested_path=None))
  
      # stable order (nested first then normal, but preserve user order otherwise)
      seen: Set[Tuple[str, Optional[str]]] = set()
      dedup: List[CoverageField] = []
      for cf in out:
          key = (cf.field, cf.nested_path)
          if key in seen:
              continue
          seen.add(key)
          dedup.append(cf)
      return dedup
  
  
  def _count_exists(es, index: str, cf: CoverageField) -> int:
      """
      Count docs where field exists.
      - If nested_path is set, uses nested query (safe for nested fields).
      - If nested query fails because path isn't actually nested in that index,
        fall back to a non-nested exists query to avoid crashing the whole report.
      """
      if cf.nested_path:
          nested_body = {
              "query": {
                  "nested": {
                      "path": cf.nested_path,
                      "query": {"exists": {"field": cf.field}},
                  }
              }
          }
          try:
              return int(es.count(index, body=nested_body))
          except Exception as e:
              # Most common: "[nested] failed to find nested object under path [...]"
              print(f"[warn] nested exists failed for {index} field={cf.field} path={cf.nested_path}: {type(e).__name__}")
              # fall through to exists
      body = {"query": {"exists": {"field": cf.field}}}
      return int(es.count(index, body=body))
  
  
  def _print_json(obj: Any) -> None:
      print(json.dumps(obj, ensure_ascii=False, indent=2, sort_keys=False))
  
  
  def compare_mapping(index_a: str, index_b: str, mapping_a: Dict[str, Any], mapping_b: Dict[str, Any]) -> None:
      flat_a = _walk_mapping_properties(mapping_a.get("mappings", {}).get("properties", {}) or {})
      flat_b = _walk_mapping_properties(mapping_b.get("mappings", {}).get("properties", {}) or {})
  
      only_a = sorted(set(flat_a) - set(flat_b))
      only_b = sorted(set(flat_b) - set(flat_a))
      type_diff = sorted([k for k in set(flat_a) & set(flat_b) if flat_a[k] != flat_b[k]])
  
      print("\n" + "=" * 90)
      print("Mapping diff (flattened field paths + types)")
      print("=" * 90)
      print(f"index_a: {index_a}")
      print(f"index_b: {index_b}")
      print(f"only_in_a: {len(only_a)}")
      print(f"only_in_b: {len(only_b)}")
      print(f"type_diff: {len(type_diff)}")
  
      if only_a[:50]:
          print("\nFields only in index_a (first 50):")
          for f in only_a[:50]:
              print(f"  - {f} ({flat_a.get(f)})")
          if len(only_a) > 50:
              print(f"  ... and {len(only_a) - 50} more")
  
      if only_b[:50]:
          print("\nFields only in index_b (first 50):")
          for f in only_b[:50]:
              print(f"  - {f} ({flat_b.get(f)})")
          if len(only_b) > 50:
              print(f"  ... and {len(only_b) - 50} more")
  
      if type_diff[:50]:
          print("\nFields with different types (first 50):")
          for f in type_diff[:50]:
              print(f"  - {f}: a={flat_a.get(f)} b={flat_b.get(f)}")
          if len(type_diff) > 50:
              print(f"  ... and {len(type_diff) - 50} more")
  
  
  def compare_coverage(
      es,
      index_a: str,
      index_b: str,
      mapping_a: Dict[str, Any],
      mapping_b: Dict[str, Any],
      fields: List[str],
      nested_fields: List[str],
  ) -> None:
      cov_fields_a = _infer_coverage_fields(mapping_a, fields, nested_fields)
      cov_fields_b = _infer_coverage_fields(mapping_b, fields, nested_fields)
  
      # keep shared list, but warn if inference differs (it shouldn't)
      if [c.field for c in cov_fields_a] != [c.field for c in cov_fields_b]:
          print("\n[warn] coverage field list differs between indices; using index_a inference as baseline")
      cov_fields = cov_fields_a
  
      print("\n" + "=" * 90)
      print("Field coverage stats (count of docs where field exists)")
      print("=" * 90)
      print(f"index_a: {index_a}")
      print(f"index_b: {index_b}")
  
      for cf in cov_fields:
          mode = f"nested(path={cf.nested_path})" if cf.nested_path else "exists"
          a = _count_exists(es, index_a, cf)
          b = _count_exists(es, index_b, cf)
          print(f"\n- {cf.field}  [{mode}]")
          print(f"  {index_a}: {a}")
          print(f"  {index_b}: {b}")
  
  
  def compare_random_samples(
      es,
      index_a: str,
      index_b: str,
      sample_size: int,
      random_seed: Optional[int],
  ) -> None:
      print("\n" + "=" * 90)
      print("Random sample diff (same _id; diff _source field paths)")
      print("=" * 90)
      print(f"sample_size: {sample_size}")
  
      random_score: Dict[str, Any] = {}
      if random_seed is not None:
          random_score["seed"] = random_seed
  
      sample_body = {
          "size": sample_size,
          "_source": False,
          "query": {"function_score": {"query": {"match_all": {}}, "random_score": random_score}},
      }
  
      # Use the underlying client directly to avoid passing duplicate `size`
      # parameters through the wrapper.
      resp = es.client.search(index=index_a, body=sample_body)
      hits = (((resp or {}).get("hits") or {}).get("hits") or [])
      ids = [h.get("_id") for h in hits if h.get("_id") is not None]
  
      if not ids:
          print("No hits returned; cannot sample.")
          return
  
      # mget in chunks
      def mget(index: str, ids_: List[str]) -> Dict[str, Dict[str, Any]]:
          out: Dict[str, Dict[str, Any]] = {}
          for batch in _chunks(ids_, 500):
              docs = es.client.mget(index=index, body={"ids": batch}).get("docs") or []
              for d in docs:
                  if d.get("found") and d.get("_id") and isinstance(d.get("_source"), dict):
                      out[d["_id"]] = d["_source"]
          return out
  
      a_docs = mget(index_a, ids)
      b_docs = mget(index_b, ids)
  
      missing_in_b = [i for i in ids if i in a_docs and i not in b_docs]
      missing_in_a = [i for i in ids if i in b_docs and i not in a_docs]
  
      only_in_a: Set[str] = set()
      only_in_b: Set[str] = set()
  
      matched = 0
      for _id in ids:
          if _id in a_docs and _id in b_docs:
              matched += 1
              pa = _field_paths_from_source(a_docs[_id])
              pb = _field_paths_from_source(b_docs[_id])
              only_in_a |= (pa - pb)
              only_in_b |= (pb - pa)
  
      summary = {
          "sample_size": len(ids),
          "matched": matched,
          "missing_in_index_b_count": len(missing_in_b),
          "missing_in_index_a_count": len(missing_in_a),
          "missing_in_index_b_example": missing_in_b[:5],
          "missing_in_index_a_example": missing_in_a[:5],
          "fields_only_in_index_a_count": len(only_in_a),
          "fields_only_in_index_b_count": len(only_in_b),
          "fields_only_in_index_a_first80": sorted(list(only_in_a))[:80],
          "fields_only_in_index_b_first80": sorted(list(only_in_b))[:80],
      }
      _print_json(summary)
  
  
  def main() -> int:
      parser = argparse.ArgumentParser(description="Compare two ES indices (mapping + data coverage + random sample).")
      parser.add_argument("index_a", help="Index A name")
      parser.add_argument("index_b", help="Index B name")
      parser.add_argument("--sample-size", type=int, default=25, help="Random sample size (default: 25)")
      parser.add_argument("--seed", type=int, default=None, help="Random seed for random_score (optional)")
      parser.add_argument(
          "--es-url",
          default=None,
          help="Elasticsearch URL. If omitted, uses env ES (preferred) or config/config.yaml.",
      )
      parser.add_argument(
          "--es-auth",
          default=None,
          help="Basic auth in 'user:pass' form. If omitted, uses env ES_AUTH or config credentials.",
      )
      parser.add_argument(
          "--fields",
          default="title.zh,vendor.zh,keywords.zh,tags.zh,keywords.en,tags.en,enriched_taxonomy_attributes,image_embedding.url,enriched_attributes.name",
          help="Comma-separated fields to compute coverage for (default: a sensible set)",
      )
      parser.add_argument(
          "--fields-nested",
          default="image_embedding.url,enriched_attributes.name",
          help="Comma-separated fields that must be treated as nested exists (default: image_embedding.url,enriched_attributes.name)",
      )
      args = parser.parse_args()
  
      # Prefer doc-style env vars (ES/ES_AUTH) to match ops workflow in docs/常用查询 - ES.md.
      # Fallback to config/config.yaml for repo-local tooling.
      env = __import__("os").environ
      es_url = args.es_url or (env.get("ES") or env.get("ES_HOST") or None)
      es_auth = args.es_auth or env.get("ES_AUTH")
      # Doc convention: if ES is unset, default to localhost:9200.
      if not es_url and es_auth:
          es_url = "http://127.0.0.1:9200"
  
      if es_url:
          username = password = None
          if es_auth and ":" in es_auth:
              username, password = es_auth.split(":", 1)
          es = ESClient(hosts=[es_url], username=username, password=password)
      else:
          es = get_es_client_from_env()
  
      if not es.ping():
          print("✗ Cannot connect to Elasticsearch")
          return 2
  
      if not es.index_exists(args.index_a):
          print(f"✗ index not found: {args.index_a}")
          return 2
      if not es.index_exists(args.index_b):
          print(f"✗ index not found: {args.index_b}")
          return 2
  
      mapping_all_a = es.get_mapping(args.index_a) or {}
      mapping_all_b = es.get_mapping(args.index_b) or {}
      if args.index_a not in mapping_all_a or args.index_b not in mapping_all_b:
          print("✗ Failed to fetch mappings for both indices")
          return 2
  
      mapping_a = mapping_all_a[args.index_a]
      mapping_b = mapping_all_b[args.index_b]
  
      compare_mapping(args.index_a, args.index_b, mapping_a, mapping_b)
  
      fields = [x for x in (args.fields or "").split(",") if x.strip()]
      nested_fields = [x for x in (args.fields_nested or "").split(",") if x.strip()]
      compare_coverage(es, args.index_a, args.index_b, mapping_a, mapping_b, fields, nested_fields)
  
      compare_random_samples(es, args.index_a, args.index_b, args.sample_size, args.seed)
  
      return 0
  
  
  if __name__ == "__main__":
      raise SystemExit(main())