#!/usr/bin/env python3
"""
Elasticsearch Index Restore Script
===================================
Restores an Elasticsearch index from a backup file created by backup_index.py

Usage:
    python restore_index.py --backup backup_file.json.gz
    python restore_index.py --backup backup_file.json.gz --target new_index_name
"""

import os
import sys
import json
import gzip
import argparse
from typing import Dict, Any, Generator

try:
    from elasticsearch import Elasticsearch
    from elasticsearch.helpers import bulk
except ImportError:
    print("ERROR: elasticsearch package not installed.")
    print("Run: pip install elasticsearch")
    sys.exit(1)

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = lambda: None


def load_config() -> Dict[str, str]:
    """Load configuration from .env file."""
    load_dotenv()
    return {
        'es_host': os.getenv('SOURCE_ES_HOST', 'http://localhost:9200'),
    }


def connect_elasticsearch(host: str) -> Elasticsearch:
    """Create Elasticsearch connection."""
    es = Elasticsearch(
        hosts=[host],
        verify_certs=False,
        request_timeout=300,
    )
    
    if not es.ping():
        raise ConnectionError(f"Cannot connect to Elasticsearch at {host}")
    
    info = es.info()
    print(f"✅ Connected to Elasticsearch {info['version']['number']}")
    return es


def read_backup_file(backup_file: str) -> Generator[Dict[str, Any], None, None]:
    """
    Read backup file and yield records.
    First record is the header with mapping/settings.
    Subsequent records are documents.
    """
    if backup_file.endswith('.gz'):
        file_handle = gzip.open(backup_file, 'rt', encoding='utf-8')
    else:
        file_handle = open(backup_file, 'r', encoding='utf-8')
    
    try:
        for line in file_handle:
            yield json.loads(line.strip())
    finally:
        file_handle.close()


def restore_index(
    es: Elasticsearch,
    backup_file: str,
    target_index: str = None,
    overwrite: bool = False,
) -> Dict[str, Any]:
    """
    Restore an Elasticsearch index from a backup file.
    """
    print(f"\n📋 Reading backup file: {backup_file}")
    
    records = read_backup_file(backup_file)
    
    # First record is the header
    header = next(records)
    if header.get('type') != 'header':
        raise ValueError("Invalid backup file: missing header")
    
    original_index = header['index']
    index_name = target_index or original_index
    
    print(f"   Original index: {original_index}")
    print(f"   Target index: {index_name}")
    print(f"   Backup timestamp: {header.get('backup_timestamp', 'unknown')}")
    print(f"   Original doc count: {header.get('doc_count', 'unknown'):,}")
    
    # Check if index exists
    if es.indices.exists(index=index_name):
        if overwrite:
            print(f"\n⚠️  Index '{index_name}' exists. Deleting...")
            es.indices.delete(index=index_name)
        else:
            raise ValueError(
                f"Index '{index_name}' already exists. "
                f"Use --overwrite to delete it, or --target to specify a different name."
            )
    
    # Create index with mapping
    print(f"\n🔧 Creating index with mapping...")
    
    # Clean up settings (remove read-only settings)
    settings = header.get('settings', {})
    index_settings = settings.get('index', {})
    
    readonly_settings = [
        'creation_date', 'provided_name', 'uuid', 'version',
        'resize', 'routing', 'verified_before_close'
    ]
    for key in readonly_settings:
        index_settings.pop(key, None)
    
    create_body = {
        'mappings': header['mapping'],
    }
    
    if index_settings:
        create_body['settings'] = {'index': index_settings}
    
    es.indices.create(index=index_name, body=create_body)
    print(f"✅ Index '{index_name}' created")
    
    # Import documents
    print(f"\n📥 Importing documents...")
    
    total_success = 0
    total_failed = 0
    batch = []
    batch_size = 500
    batch_num = 0
    
    for record in records:
        if record.get('type') != 'document':
            continue
            
        action = {
            '_index': index_name,
            '_id': record['_id'],
            '_source': record['_source'],
        }
        batch.append(action)
        
        if len(batch) >= batch_size:
            batch_num += 1
            success, failed = bulk(
                es,
                batch,
                raise_on_error=False,
                raise_on_exception=False,
            )
            total_success += success
            total_failed += len(failed) if isinstance(failed, list) else 0
            print(f"   Batch {batch_num}: {total_success:,} docs imported...", end='\r')
            batch = []
    
    # Process remaining batch
    if batch:
        batch_num += 1
        success, failed = bulk(
            es,
            batch,
            raise_on_error=False,
            raise_on_exception=False,
        )
        total_success += success
        total_failed += len(failed) if isinstance(failed, list) else 0
    
    print(f"\n✅ Import complete")
    print(f"   Success: {total_success:,}")
    print(f"   Failed: {total_failed:,}")
    
    # Refresh index
    es.indices.refresh(index=index_name)
    
    # Verify
    stats = es.indices.stats(index=index_name)
    final_count = stats['indices'][index_name]['primaries']['docs']['count']
    print(f"\n📊 Final document count: {final_count:,}")
    
    return {
        'index': index_name,
        'success': total_success,
        'failed': total_failed,
        'final_count': final_count,
    }


def main():
    parser = argparse.ArgumentParser(description='Restore Elasticsearch index from backup')
    parser.add_argument('--backup', '-b', required=True, help='Path to backup file (.json.gz)')
    parser.add_argument('--target', '-t', help='Target index name (default: original name)')
    parser.add_argument('--host', '-H', help='Elasticsearch host URL')
    parser.add_argument('--overwrite', action='store_true', help='Overwrite if index exists')
    
    args = parser.parse_args()
    
    config = load_config()
    es_host = args.host or config['es_host']
    
    print(f"""
╔══════════════════════════════════════════════════════════════════╗
║           ELASTICSEARCH INDEX RESTORE                             ║
╚══════════════════════════════════════════════════════════════════╝
    """)
    
    print(f"🔧 Configuration:")
    print(f"   Host: {es_host.split('@')[-1] if '@' in es_host else es_host}")
    print(f"   Backup: {args.backup}")
    if args.target:
        print(f"   Target: {args.target}")
    
    try:
        es = connect_elasticsearch(es_host)
        
        result = restore_index(
            es,
            args.backup,
            target_index=args.target,
            overwrite=args.overwrite,
        )
        
        print("\n🎉 Restore completed successfully!")
        return 0
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        return 1


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