"""
Milvus Schema Migration Script for Domain-Aware RAG.

Adds domain_category and domain_terms fields to existing Milvus collection.
This script performs a zero-downtime migration by:
1. Creating a new collection with updated schema
2. Copying existing data to new collection
3. Replacing old collection with new one

IMPORTANT: Run this during low-traffic period. Backup data before running.

Usage:
    python scripts/migrate_milvus_schema.py [--dry-run]
"""

import argparse
import sys
import time
from typing import Any

import structlog
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    connections,
    utility,
)

from src.core.config import settings

logger = structlog.get_logger(__name__)


def create_updated_schema(vector_dim: int = 3072) -> CollectionSchema:
    """Create updated Milvus schema with domain fields.

    Args:
        vector_dim: Dimension of embedding vectors (default: 3072 for text-embedding-3-large)

    Returns:
        CollectionSchema with domain_category and domain_terms fields added
    """
    fields = [
        FieldSchema(
            name="id",
            dtype=DataType.INT64,
            is_primary=True,
            auto_id=True,
        ),
        FieldSchema(
            name="vector",
            dtype=DataType.FLOAT_VECTOR,
            dim=vector_dim,
        ),
        FieldSchema(
            name="record_id",
            dtype=DataType.VARCHAR,
            max_length=36,
        ),
        FieldSchema(
            name="document_id",
            dtype=DataType.VARCHAR,
            max_length=36,
        ),
        FieldSchema(
            name="filename",
            dtype=DataType.VARCHAR,
            max_length=255,
        ),
        FieldSchema(
            name="document_type",
            dtype=DataType.VARCHAR,
            max_length=20,
        ),
        FieldSchema(
            name="text_content",
            dtype=DataType.VARCHAR,
            max_length=65535,
        ),
        FieldSchema(
            name="metadata",
            dtype=DataType.JSON,
        ),
        # NEW FIELDS for Domain-Aware RAG
        FieldSchema(
            name="domain_category",
            dtype=DataType.VARCHAR,
            max_length=64,
            default_value="general",  # Default for backward compatibility
        ),
        FieldSchema(
            name="domain_terms",
            dtype=DataType.ARRAY,
            element_type=DataType.VARCHAR,
            max_capacity=50,  # Max 50 terms per chunk
            max_length=100,  # Max 100 chars per term
        ),
    ]

    schema = CollectionSchema(
        fields,
        description="DSOL document records with embeddings and domain terms",
    )

    return schema


def migrate_collection(
    old_collection_name: str,
    dry_run: bool = False
) -> None:
    """Migrate existing collection to new schema with domain fields.

    Migration steps:
    1. Verify old collection exists
    2. Create temporary new collection with updated schema
    3. Copy all data from old to new collection
    4. Create indexes on new collection
    5. Drop old collection
    6. Rename new collection to original name

    Args:
        old_collection_name: Name of existing collection to migrate
        dry_run: If True, only simulate migration without making changes

    Raises:
        ValueError: If old collection doesn't exist
        Exception: If migration fails
    """
    logger.info(
        "migration_started",
        collection=old_collection_name,
        dry_run=dry_run
    )

    # Connect to Milvus
    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )

    # Verify old collection exists
    if not utility.has_collection(old_collection_name):
        raise ValueError(f"Collection '{old_collection_name}' does not exist")

    # Get old collection info
    old_collection = Collection(old_collection_name)
    old_collection.load()

    num_entities = old_collection.num_entities
    logger.info("old_collection_info", num_entities=num_entities)

    if dry_run:
        logger.info("dry_run_mode", message="Would migrate collection", entities=num_entities)
        return

    # Create new collection with updated schema
    new_collection_name = f"{old_collection_name}_migrated"
    logger.info("creating_new_collection", name=new_collection_name)

    # Drop if exists (in case of retry)
    if utility.has_collection(new_collection_name):
        logger.warning("dropping_existing_temp_collection", name=new_collection_name)
        utility.drop_collection(new_collection_name)

    # Get vector dimension from old collection
    vector_field = None
    for field in old_collection.schema.fields:
        if field.dtype == DataType.FLOAT_VECTOR:
            vector_field = field
            break

    if not vector_field:
        raise ValueError("No vector field found in old collection")

    vector_dim = vector_field.params.get("dim", 3072)
    logger.info("detected_vector_dim", dim=vector_dim)

    # Create new collection
    new_schema = create_updated_schema(vector_dim)
    new_collection = Collection(
        name=new_collection_name,
        schema=new_schema,
    )

    # Copy data in batches
    logger.info("copying_data", total_entities=num_entities)

    batch_size = 1000
    copied = 0

    try:
        # Query all data from old collection in batches
        for offset in range(0, num_entities, batch_size):
            # Query batch from old collection
            results = old_collection.query(
                expr=f"id >= 0",  # Get all entities
                output_fields=["*"],
                limit=batch_size,
                offset=offset,
            )

            if not results:
                break

            # Prepare data for new collection
            # Add empty arrays for domain_terms and "general" for domain_category
            batch_data = []
            for entity in results:
                # Preserve existing fields
                new_entity = [
                    entity["vector"],
                    entity["record_id"],
                    entity["document_id"],
                    entity["filename"],
                    entity["document_type"],
                    entity["text_content"],
                    entity["metadata"],
                    "general",  # domain_category default
                    [],  # domain_terms empty array
                ]
                batch_data.append(new_entity)

            # Transpose for columnar format
            columnar_data = list(map(list, zip(*batch_data)))

            # Insert into new collection
            new_collection.insert(columnar_data)
            new_collection.flush()

            copied += len(results)
            logger.info(
                "batch_copied",
                offset=offset,
                batch_size=len(results),
                total_copied=copied,
                progress_pct=round((copied / num_entities) * 100, 1)
            )

            # Small delay to avoid overwhelming Milvus
            time.sleep(0.1)

        logger.info("data_copy_complete", total_copied=copied)

        # Create indexes on new collection
        logger.info("creating_indexes")

        # Vector index (HNSW for similarity search)
        vector_index_params = {
            "metric_type": "COSINE",
            "index_type": "HNSW",
            "params": {
                "M": 16,
                "efConstruction": 200,
            },
        }
        new_collection.create_index(
            field_name="vector",
            index_params=vector_index_params,
        )

        # Scalar index for domain_category (for filtering)
        new_collection.create_index(
            field_name="domain_category",
            index_params={"index_type": "INVERTED"},
        )

        # Array index for domain_terms (for array_contains_any filtering)
        new_collection.create_index(
            field_name="domain_terms",
            index_params={"index_type": "INVERTED"},
        )

        logger.info("indexes_created")

        # Load new collection into memory
        new_collection.load()
        logger.info("new_collection_loaded")

        # Verify data integrity
        new_num_entities = new_collection.num_entities
        if new_num_entities != num_entities:
            raise ValueError(
                f"Data integrity check failed: "
                f"old={num_entities}, new={new_num_entities}"
            )

        logger.info("data_integrity_verified", entities=new_num_entities)

        # Release old collection
        old_collection.release()

        # Drop old collection
        logger.info("dropping_old_collection", name=old_collection_name)
        utility.drop_collection(old_collection_name)

        # Rename new collection to old name
        logger.info("renaming_new_collection", old_name=new_collection_name, new_name=old_collection_name)
        utility.rename_collection(new_collection_name, old_collection_name)

        logger.info(
            "migration_complete",
            collection=old_collection_name,
            entities_migrated=new_num_entities
        )

    except Exception as e:
        logger.error(
            "migration_failed",
            error=str(e),
            copied=copied,
            total=num_entities
        )
        # Clean up new collection on failure
        if utility.has_collection(new_collection_name):
            logger.info("cleaning_up_failed_migration", collection=new_collection_name)
            utility.drop_collection(new_collection_name)
        raise

    finally:
        # Disconnect
        connections.disconnect("default")


def verify_schema(collection_name: str) -> None:
    """Verify that collection has updated schema with domain fields.

    Args:
        collection_name: Name of collection to verify
    """
    logger.info("verifying_schema", collection=collection_name)

    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )

    if not utility.has_collection(collection_name):
        logger.error("collection_not_found", name=collection_name)
        return

    collection = Collection(collection_name)

    # Check for domain fields
    field_names = [field.name for field in collection.schema.fields]

    has_category = "domain_category" in field_names
    has_terms = "domain_terms" in field_names

    logger.info(
        "schema_verification",
        has_domain_category=has_category,
        has_domain_terms=has_terms,
        all_fields=field_names,
    )

    if has_category and has_terms:
        logger.info("schema_verification_passed")
    else:
        logger.warning("schema_verification_failed", missing_fields=[
            f for f in ["domain_category", "domain_terms"]
            if f not in field_names
        ])

    connections.disconnect("default")


def main():
    """Main entry point for migration script."""
    parser = argparse.ArgumentParser(
        description="Migrate Milvus collection to add domain fields"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Simulate migration without making changes"
    )
    parser.add_argument(
        "--verify-only",
        action="store_true",
        help="Only verify schema, don't migrate"
    )
    parser.add_argument(
        "--collection",
        default=settings.milvus_collection,
        help=f"Collection name (default: {settings.milvus_collection})"
    )

    args = parser.parse_args()

    try:
        if args.verify_only:
            verify_schema(args.collection)
        else:
            migrate_collection(args.collection, dry_run=args.dry_run)

    except Exception as e:
        logger.error("migration_error", error=str(e))
        sys.exit(1)


if __name__ == "__main__":
    main()
