"""
  Data Indexing Script
  Generates embeddings for products and stores them in Milvus
  """
  
  import csv
  import logging
  import os
  import sys
  from pathlib import Path
  from typing import Any, Dict, Optional
  
  from tqdm import tqdm
  
  # Add parent directory to path
  sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  
  # Import config and settings first
  # Direct imports from files to avoid __init__.py circular issues
  import importlib.util
  
  from app.config import get_absolute_path, settings
  
  
  def load_service_module(module_name, file_name):
      """Load a service module directly from file"""
      spec = importlib.util.spec_from_file_location(
          module_name,
          os.path.join(
              os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
              f"app/services/{file_name}",
          ),
      )
      module = importlib.util.module_from_spec(spec)
      spec.loader.exec_module(module)
      return module
  
  
  embedding_module = load_service_module("embedding_service", "embedding_service.py")
  milvus_module = load_service_module("milvus_service", "milvus_service.py")
  
  EmbeddingService = embedding_module.EmbeddingService
  MilvusService = milvus_module.MilvusService
  
  # Configure logging
  logging.basicConfig(
      level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
  )
  logger = logging.getLogger(__name__)
  
  
  class DataIndexer:
      """Index product data by generating and storing embeddings"""
  
      def __init__(self):
          """Initialize services"""
          self.embedding_service = EmbeddingService()
          self.milvus_service = MilvusService()
  
          self.image_dir = Path(get_absolute_path(settings.image_data_path))
          self.styles_csv = get_absolute_path("./data/styles.csv")
          self.images_csv = get_absolute_path("./data/images.csv")
  
          # Load product data from CSV
          self.products = self._load_products_from_csv()
  
      def _load_products_from_csv(self) -> Dict[int, Dict[str, Any]]:
          """Load products from CSV files"""
          products = {}
  
          # Load images mapping
          images_dict = {}
          with open(self.images_csv, "r", encoding="utf-8") as f:
              reader = csv.DictReader(f)
              for row in reader:
                  product_id = int(row["filename"].split(".")[0])
                  images_dict[product_id] = row["link"]
  
          # Load styles/products
          with open(self.styles_csv, "r", encoding="utf-8") as f:
              reader = csv.DictReader(f)
              for row in reader:
                  try:
                      product_id = int(row["id"])
                      products[product_id] = {
                          "id": product_id,
                          "gender": row.get("gender", ""),
                          "masterCategory": row.get("masterCategory", ""),
                          "subCategory": row.get("subCategory", ""),
                          "articleType": row.get("articleType", ""),
                          "baseColour": row.get("baseColour", ""),
                          "season": row.get("season", ""),
                          "year": int(row["year"]) if row.get("year") else 0,
                          "usage": row.get("usage", ""),
                          "productDisplayName": row.get("productDisplayName", ""),
                          "imageUrl": images_dict.get(product_id, ""),
                          "imagePath": f"{product_id}.jpg",
                      }
                  except (ValueError, KeyError) as e:
                      logger.warning(f"Error loading product {row.get('id')}: {e}")
                      continue
  
          logger.info(f"Loaded {len(products)} products from CSV")
          return products
  
      def setup(self) -> None:
          """Setup connections and collections"""
          logger.info("Setting up services...")
  
          # Connect to CLIP server
          self.embedding_service.connect_clip()
          logger.info("✓ CLIP server connected")
  
          # Connect to Milvus
          self.milvus_service.connect()
          logger.info("✓ Milvus connected")
  
          # Create Milvus collections
          self.milvus_service.create_text_collection(recreate=False)
          self.milvus_service.create_image_collection(recreate=False)
          logger.info("✓ Milvus collections ready")
  
      def teardown(self) -> None:
          """Close all connections"""
          logger.info("Closing connections...")
          self.embedding_service.disconnect_clip()
          self.milvus_service.disconnect()
          logger.info("✓ All connections closed")
  
      def index_text_embeddings(
          self, batch_size: int = 100, skip: int = 0, limit: Optional[int] = None
      ) -> Dict[str, int]:
          """Generate and store text embeddings for products
  
          Args:
              batch_size: Number of products to process at once
              skip: Number of products to skip
              limit: Maximum number of products to process (None for all)
  
          Returns:
              Dictionary with indexing statistics
          """
          logger.info("Starting text embedding indexing...")
  
          # Get products list
          product_ids = list(self.products.keys())[skip:]
          if limit:
              product_ids = product_ids[:limit]
  
          total_products = len(product_ids)
          processed = 0
          inserted = 0
          errors = 0
  
          with tqdm(total=total_products, desc="Indexing text embeddings") as pbar:
              while processed < total_products:
                  # Get batch of products
                  current_batch_size = min(batch_size, total_products - processed)
                  batch_ids = product_ids[processed : processed + current_batch_size]
                  products = [self.products[pid] for pid in batch_ids]
  
                  if not products:
                      break
  
                  try:
                      # Prepare texts for embedding
                      texts = []
                      text_mappings = []
  
                      for product in products:
                          # Create text representation of product
                          text = self._create_product_text(product)
                          texts.append(text)
                          text_mappings.append(
                              {"product_id": product["id"], "text": text}
                          )
  
                      # Generate embeddings
                      embeddings = self.embedding_service.get_text_embeddings_batch(
                          texts, batch_size=50  # OpenAI batch size
                      )
  
                      # Prepare data for Milvus (with metadata)
                      milvus_data = []
                      for idx, (mapping, embedding) in enumerate(
                          zip(text_mappings, embeddings)
                      ):
                          product_id = mapping["product_id"]
                          product = self.products[product_id]
  
                          milvus_data.append(
                              {
                                  "id": product_id,
                                  "text": mapping["text"][
                                      :2000
                                  ],  # Truncate to max length
                                  "embedding": embedding,
                                  # Product metadata
                                  "productDisplayName": product["productDisplayName"][
                                      :500
                                  ],
                                  "gender": product["gender"][:50],
                                  "masterCategory": product["masterCategory"][:100],
                                  "subCategory": product["subCategory"][:100],
                                  "articleType": product["articleType"][:100],
                                  "baseColour": product["baseColour"][:50],
                                  "season": product["season"][:50],
                                  "usage": product["usage"][:50],
                                  "year": product["year"],
                                  "imageUrl": product["imageUrl"],
                                  "imagePath": product["imagePath"],
                              }
                          )
  
                      # Insert into Milvus
                      count = self.milvus_service.insert_text_embeddings(milvus_data)
                      inserted += count
  
                  except Exception as e:
                      logger.error(
                          f"Error processing text batch at offset {processed}: {e}"
                      )
                      errors += len(products)
  
                  processed += len(products)
                  pbar.update(len(products))
  
          stats = {"total_processed": processed, "inserted": inserted, "errors": errors}
  
          logger.info(f"Text embedding indexing completed: {stats}")
          return stats
  
      def index_image_embeddings(
          self, batch_size: int = 32, skip: int = 0, limit: Optional[int] = None
      ) -> Dict[str, int]:
          """Generate and store image embeddings for products
  
          Args:
              batch_size: Number of images to process at once
              skip: Number of products to skip
              limit: Maximum number of products to process (None for all)
  
          Returns:
              Dictionary with indexing statistics
          """
          logger.info("Starting image embedding indexing...")
  
          # Get products list
          product_ids = list(self.products.keys())[skip:]
          if limit:
              product_ids = product_ids[:limit]
  
          total_products = len(product_ids)
          processed = 0
          inserted = 0
          errors = 0
  
          with tqdm(total=total_products, desc="Indexing image embeddings") as pbar:
              while processed < total_products:
                  # Get batch of products
                  current_batch_size = min(batch_size, total_products - processed)
                  batch_ids = product_ids[processed : processed + current_batch_size]
                  products = [self.products[pid] for pid in batch_ids]
  
                  if not products:
                      break
  
                  try:
                      # Prepare image paths
                      image_paths = []
                      image_mappings = []
  
                      for product in products:
                          image_path = self.image_dir / product["imagePath"]
                          image_paths.append(image_path)
                          image_mappings.append(
                              {
                                  "product_id": product["id"],
                                  "image_path": product["imagePath"],
                              }
                          )
  
                      # Generate embeddings
                      embeddings = self.embedding_service.get_image_embeddings_batch(
                          image_paths, batch_size=batch_size
                      )
  
                      # Prepare data for Milvus (with metadata)
                      milvus_data = []
                      for idx, (mapping, embedding) in enumerate(
                          zip(image_mappings, embeddings)
                      ):
                          if embedding is not None:
                              product_id = mapping["product_id"]
                              product = self.products[product_id]
  
                              milvus_data.append(
                                  {
                                      "id": product_id,
                                      "image_path": mapping["image_path"],
                                      "embedding": embedding,
                                      # Product metadata
                                      "productDisplayName": product["productDisplayName"][
                                          :500
                                      ],
                                      "gender": product["gender"][:50],
                                      "masterCategory": product["masterCategory"][:100],
                                      "subCategory": product["subCategory"][:100],
                                      "articleType": product["articleType"][:100],
                                      "baseColour": product["baseColour"][:50],
                                      "season": product["season"][:50],
                                      "usage": product["usage"][:50],
                                      "year": product["year"],
                                      "imageUrl": product["imageUrl"],
                                  }
                              )
                          else:
                              errors += 1
  
                      # Insert into Milvus
                      if milvus_data:
                          count = self.milvus_service.insert_image_embeddings(milvus_data)
                          inserted += count
  
                  except Exception as e:
                      logger.error(
                          f"Error processing image batch at offset {processed}: {e}"
                      )
                      errors += len(products)
  
                  processed += len(products)
                  pbar.update(len(products))
  
          stats = {"total_processed": processed, "inserted": inserted, "errors": errors}
  
          logger.info(f"Image embedding indexing completed: {stats}")
          return stats
  
      def _create_product_text(self, product: Dict[str, Any]) -> str:
          """Create text representation of product for embedding
  
          Args:
              product: Product document
  
          Returns:
              Text representation
          """
          # Create a natural language description
          parts = [
              product.get("productDisplayName", ""),
              f"Gender: {product.get('gender', '')}",
              f"Category: {product.get('masterCategory', '')} > {product.get('subCategory', '')}",
              f"Type: {product.get('articleType', '')}",
              f"Color: {product.get('baseColour', '')}",
              f"Season: {product.get('season', '')}",
              f"Usage: {product.get('usage', '')}",
          ]
  
          text = " | ".join(
              [p for p in parts if p and p != "Gender: " and p != "Color: "]
          )
          return text
  
      def get_stats(self) -> Dict[str, Any]:
          """Get indexing statistics
  
          Returns:
              Dictionary with statistics
          """
          text_stats = self.milvus_service.get_collection_stats(
              self.milvus_service.text_collection_name
          )
          image_stats = self.milvus_service.get_collection_stats(
              self.milvus_service.image_collection_name
          )
  
          return {
              "total_products": len(self.products),
              "milvus_text": text_stats,
              "milvus_image": image_stats,
          }
  
  
  def main():
      """Main function"""
      import argparse
  
      parser = argparse.ArgumentParser(description="Index product data for search")
      parser.add_argument(
          "--mode",
          choices=["text", "image", "both"],
          default="both",
          help="Which embeddings to index",
      )
      parser.add_argument(
          "--batch-size", type=int, default=100, help="Batch size for processing"
      )
      parser.add_argument(
          "--skip", type=int, default=0, help="Number of products to skip"
      )
      parser.add_argument(
          "--limit", type=int, default=None, help="Maximum number of products to process"
      )
      parser.add_argument("--stats", action="store_true", help="Show statistics only")
  
      args = parser.parse_args()
  
      # Create indexer
      indexer = DataIndexer()
  
      try:
          # Setup services
          indexer.setup()
  
          if args.stats:
              # Show statistics
              stats = indexer.get_stats()
              print("\n=== Indexing Statistics ===")
              print(f"\nTotal Products in CSV: {stats['total_products']}")
  
              print("\nMilvus Text Embeddings:")
              print(f"  Collection: {stats['milvus_text']['collection_name']}")
              print(f"  Total embeddings: {stats['milvus_text']['row_count']}")
  
              print("\nMilvus Image Embeddings:")
              print(f"  Collection: {stats['milvus_image']['collection_name']}")
              print(f"  Total embeddings: {stats['milvus_image']['row_count']}")
  
              print(
                  f"\nCoverage: {stats['milvus_image']['row_count'] / stats['total_products'] * 100:.1f}%"
              )
          else:
              # Index data
              if args.mode in ["text", "both"]:
                  logger.info("=== Indexing Text Embeddings ===")
                  text_stats = indexer.index_text_embeddings(
                      batch_size=args.batch_size, skip=args.skip, limit=args.limit
                  )
                  print(f"\nText Indexing Results: {text_stats}")
  
              if args.mode in ["image", "both"]:
                  logger.info("=== Indexing Image Embeddings ===")
                  image_stats = indexer.index_image_embeddings(
                      batch_size=min(args.batch_size, 32),  # Smaller batch for images
                      skip=args.skip,
                      limit=args.limit,
                  )
                  print(f"\nImage Indexing Results: {image_stats}")
  
              # Show final statistics
              logger.info("\n=== Final Statistics ===")
              stats = indexer.get_stats()
              print(f"Total products: {stats['total_products']}")
              print(f"Text embeddings: {stats['milvus_text']['row_count']}")
              print(f"Image embeddings: {stats['milvus_image']['row_count']}")
  
      except KeyboardInterrupt:
          logger.info("\nIndexing interrupted by user")
      except Exception as e:
          logger.error(f"Error during indexing: {e}", exc_info=True)
          sys.exit(1)
      finally:
          indexer.teardown()
  
  
  if __name__ == "__main__":
      main()