#!/usr/bin/env python3
  """
  Customer1 data ingestion script.
  
  Loads data from CSV and indexes into Elasticsearch with embeddings.
  """
  
  import sys
  import os
  import pandas as pd
  import argparse

  from typing import Optional

  # Add parent directory to path (go up 3 levels: customer1 -> data -> SearchEngine -> root)
  project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  sys.path.insert(0, project_root)

  
  from config import ConfigLoader
  from utils import ESClient, get_connection_from_config
  from indexer import DataTransformer, IndexingPipeline
  from embeddings import BgeEncoder, CLIPImageEncoder
  
  
  def load_csv_data(csv_path: str, limit: Optional[int] = None) -> pd.DataFrame:
      """
      Load data from CSV file.
  
      Args:
          csv_path: Path to CSV file
          limit: Maximum number of rows to load (None for all)
  
      Returns:
          DataFrame with data
      """
      print(f"[Ingestion] Loading data from: {csv_path}")
  
      df = pd.read_csv(csv_path)
  
      if limit:
          df = df.head(limit)
  
      print(f"[Ingestion] Loaded {len(df)} rows")
      print(f"[Ingestion] Columns: {df.columns.tolist()}")
  
      return df
  
  
  def main():
      """Main ingestion function."""
      parser = argparse.ArgumentParser(description='Ingest customer1 data into Elasticsearch')
      parser.add_argument('--config', default='customer1', help='Customer config name')
      parser.add_argument('--csv', default='data/customer1/goods_with_pic.5years_congku.csv.shuf.1w',
                         help='Path to CSV data file')
      parser.add_argument('--limit', type=int, help='Limit number of documents to index')
      parser.add_argument('--batch-size', type=int, default=100, help='Batch size for processing')
      parser.add_argument('--recreate-index', action='store_true', help='Recreate index if exists')
      parser.add_argument('--es-host', default='http://localhost:9200', help='Elasticsearch host')

      parser.add_argument('--es-username', default=None, help='Elasticsearch username (or set ES_USERNAME env var)')
      parser.add_argument('--es-password', default=None, help='Elasticsearch password (or set ES_PASSWORD env var)')

      parser.add_argument('--skip-embeddings', action='store_true', help='Skip embedding generation')
      args = parser.parse_args()
  
      print("="*60)
      print("Customer1 Data Ingestion")
      print("="*60)
  
      # Load configuration
      print(f"\n[1/6] Loading configuration: {args.config}")
      config_loader = ConfigLoader("config/schema")
      config = config_loader.load_customer_config(args.config)
  
      # Validate configuration
      errors = config_loader.validate_config(config)
      if errors:
          print(f"Configuration validation failed:")
          for error in errors:
              print(f"  - {error}")
          return 1
  
      print(f"Configuration loaded successfully")
      print(f"  - Index: {config.es_index_name}")
      print(f"  - Fields: {len(config.fields)}")
      print(f"  - Indexes: {len(config.indexes)}")
  
      # Initialize Elasticsearch client
      print(f"\n[2/6] Connecting to Elasticsearch: {args.es_host}")

      
      # Get credentials: prioritize command-line args, then environment variables, then .env file
      es_username = args.es_username
      es_password = args.es_password
      
      # If not provided via args, try to load from .env file via env_config
      if not es_username or not es_password:
          try:
              from config.env_config import get_es_config
              es_config = get_es_config()
              es_username = es_username or es_config.get('username')
              es_password = es_password or es_config.get('password')
          except Exception:
              # Fallback to environment variables
              es_username = es_username or os.getenv('ES_USERNAME')
              es_password = es_password or os.getenv('ES_PASSWORD')
      
      # Create ES client with credentials if available
      if es_username and es_password:
          print(f"  Using authentication: {es_username}")
          es_client = ESClient(hosts=[args.es_host], username=es_username, password=es_password)
      else:
          print(f"  Warning: No authentication credentials found")
          print(f"  Attempting connection without authentication (will fail if ES requires auth)")
          es_client = ESClient(hosts=[args.es_host])

  
      if not es_client.ping():
          print("Failed to connect to Elasticsearch")

          print("\nTroubleshooting:")
          print("  1. Check if Elasticsearch is running: curl http://localhost:9200")
          print("  2. If ES requires authentication, provide credentials:")
          print("     - Use --es-username and --es-password arguments, or")
          print("     - Set ES_USERNAME and ES_PASSWORD environment variables")
          print("  3. Verify the host URL is correct: --es-host")

          return 1
  
      print("Connected to Elasticsearch successfully")
  
      # Load data
      print(f"\n[3/6] Loading data from CSV")
      df = load_csv_data(args.csv, limit=args.limit)
  
      # Initialize encoders (if not skipping embeddings)
      text_encoder = None
      image_encoder = None
  
      if not args.skip_embeddings:
          print(f"\n[4/6] Initializing embedding encoders")
          print("This may take a few minutes on first run (downloading models)...")
  
          try:
              text_encoder = BgeEncoder()
              print("Text encoder initialized")
          except Exception as e:
              print(f"Warning: Failed to initialize text encoder: {e}")
              print("Continuing without text embeddings...")
  
          try:
              image_encoder = CLIPImageEncoder()
              print("Image encoder initialized")
          except Exception as e:
              print(f"Warning: Failed to initialize image encoder: {e}")
              print("Continuing without image embeddings...")
      else:
          print(f"\n[4/6] Skipping embedding generation (--skip-embeddings)")
  
      # Initialize data transformer
      print(f"\n[5/6] Initializing data transformation pipeline")
      transformer = DataTransformer(
          config=config,
          text_encoder=text_encoder,
          image_encoder=image_encoder,
          use_cache=True
      )
  
      # Run indexing pipeline
      print(f"\n[6/6] Starting indexing pipeline")
      pipeline = IndexingPipeline(
          config=config,
          es_client=es_client,
          data_transformer=transformer,
          recreate_index=args.recreate_index
      )
  
      results = pipeline.run(df, batch_size=args.batch_size)
  
      # Print summary
      print("\n" + "="*60)
      print("Ingestion Complete!")
      print("="*60)
      print(f"Total documents: {results['total']}")
      print(f"Successfully indexed: {results['success']}")
      print(f"Failed: {results['failed']}")
      print(f"Time elapsed: {results['elapsed_time']:.2f}s")
      print(f"Throughput: {results['docs_per_second']:.2f} docs/s")
  
      if results['errors']:
          print(f"\nFirst few errors:")
          for error in results['errors'][:5]:
              print(f"  - {error}")
  
      # Verify index
      print(f"\nVerifying index...")
      doc_count = es_client.count(config.es_index_name)
      print(f"Documents in index: {doc_count}")
  
      print("\nIngestion completed successfully!")
      return 0
  
  
  if __name__ == "__main__":
      sys.exit(main())