"""Reset old data and recreate indices with new mappings.

WARNING: This script DELETES all existing document data!
- Drops and recreates OpenSearch indices (gov_doc_chunks, gov_doc_meta)
- Clears all Neo4j graph data

Usage:
    cd backend
    python -m scripts.reset_and_migrate
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path

# Add backend to path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from opensearchpy import AsyncOpenSearch

from app.config import settings
from app.infrastructure.es_client import (
    ESClient,
    GOV_DOC_CHUNKS_MAPPING,
    GOV_DOC_META_MAPPING,
    HYBRID_RRF_PIPELINE,
    _RRF_PIPELINE_BODY,
)


async def reset_opensearch() -> None:
    """Delete and recreate OpenSearch indices."""
    es_kwargs: dict = {"hosts": [settings.es_host], "timeout": 60}
    if settings.es_username:
        es_kwargs["http_auth"] = (settings.es_username, settings.es_password)

    es = AsyncOpenSearch(**es_kwargs)

    try:
        for index_name in [settings.es_chunk_index, settings.es_meta_index]:
            exists = await es.indices.exists(index=index_name)
            if exists:
                print(f"  Deleting index: {index_name}")
                await es.indices.delete(index=index_name)
                print(f"  Deleted: {index_name}")

        print(f"  Creating index: {settings.es_chunk_index}")
        await es.indices.create(
            index=settings.es_chunk_index, body=GOV_DOC_CHUNKS_MAPPING
        )
        print(f"  Created: {settings.es_chunk_index}")

        print(f"  Creating index: {settings.es_meta_index}")
        await es.indices.create(
            index=settings.es_meta_index, body=GOV_DOC_META_MAPPING
        )
        print(f"  Created: {settings.es_meta_index}")

        # Recreate RRF pipeline
        try:
            await es.http.put(
                f"/_search/pipeline/{HYBRID_RRF_PIPELINE}",
                body=_RRF_PIPELINE_BODY,
            )
            print("  RRF pipeline created")
        except Exception as exc:
            print(f"  RRF pipeline warning: {exc}")

    finally:
        await es.close()


async def reset_neo4j() -> None:
    """Clear all Neo4j graph data."""
    from neo4j import AsyncGraphDatabase

    driver = AsyncGraphDatabase.driver(
        settings.neo4j_uri,
        auth=(settings.neo4j_user, settings.neo4j_password),
    )
    try:
        async with driver.session(database=settings.neo4j_database) as session:
            result = await session.run("MATCH (n) RETURN count(n) AS cnt")
            record = await result.single()
            count = record["cnt"] if record else 0
            print(f"  Found {count} nodes in Neo4j")

            if count > 0:
                print("  Deleting all nodes and relationships...")
                # Delete in batches to avoid memory issues
                while True:
                    result = await session.run(
                        "MATCH (n) WITH n LIMIT 10000 DETACH DELETE n RETURN count(*) AS deleted"
                    )
                    record = await result.single()
                    deleted = record["deleted"] if record else 0
                    if deleted == 0:
                        break
                    print(f"  Deleted {deleted} nodes")

            print("  Neo4j cleared")
    finally:
        await driver.close()


async def main() -> None:
    print("=" * 60)
    print("  ZM-RAG: Reset & Migrate")
    print("  WARNING: This will DELETE all existing document data!")
    print("=" * 60)

    confirm = input("\nType 'yes' to confirm: ").strip().lower()
    if confirm != "yes":
        print("Aborted.")
        return

    print("\n[1/2] Resetting OpenSearch indices...")
    await reset_opensearch()

    print("\n[2/2] Resetting Neo4j graph data...")
    await reset_neo4j()

    print("\n" + "=" * 60)
    print("  Done! All data has been cleared and indices recreated.")
    print("  You can now restart the backend and ingest new documents.")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(main())