"""Upgrading Existing documents to default catalog

Revision ID: 28984b05d396
Revises: 0f5932ff3051
Create Date: 2024-10-15 09:02:23.355660

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.orm import Session
import pgvector
from common.models.document import Catalog, Document
from common.models.user import Tenant
from flask import current_app
import logging
import os

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create a console handler
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# Create a logging format
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)

# Add the handler to the logger
logger.addHandler(console_handler)

logger.info("Starting revision upgrade 28984b05d396")

# revision identifiers, used by Alembic.
revision = '28984b05d396'
down_revision = '0f5932ff3051'
branch_labels = None
depends_on = None


def upgrade():
    logger.info("Starting migration: Creating default catalog and updating documents.")

    bind = op.get_bind()
    session = Session(bind=bind)

    try:
        tenant_id = int(os.getenv('TENANT_ID'))
        logger.info(f"In migration: tenant_id = {tenant_id}")

        # Step 1: Create a new Default Catalog
        logger.info("Creating default catalog")
        default_catalog = Catalog(
            name='Default Catalog',
            description=None,
        )
        tenant = Tenant.query.get_or_404(tenant_id)
        default_catalog.html_tags = tenant.html_tags
        default_catalog.html_end_tags = tenant.html_end_tags
        default_catalog.html_included_elements = tenant.html_included_elements
        default_catalog.html_excluded_elements = tenant.html_excluded_elements
        default_catalog.html_excluded_classes = tenant.html_excluded_classes
        default_catalog.min_chunk_size = tenant.min_chunk_size
        default_catalog.max_chunk_size = tenant.max_chunk_size
        default_catalog.es_k = tenant.es_k
        default_catalog.es_similarity_threshold = tenant.es_similarity_threshold
        default_catalog.chat_RAG_temperature = tenant.chat_RAG_temperature
        default_catalog.chat_no_RAG_temperature = tenant.chat_no_RAG_temperature
        default_catalog.embed_tuning = tenant.embed_tuning
        default_catalog.rag_tuning = tenant.rag_tuning

        session.add(default_catalog)
        session.commit()

        new_catalog_id = default_catalog.id
        logger.info(f"Default catalog created with ID: {new_catalog_id}")

        # Step 2: Update all documents to use this new catalog
        logger.info("Updating documents with the new catalog ID.")
        documents = session.query(Document).all()
        for document in documents:
            document.catalog_id = new_catalog_id
        session.commit()

        logger.info(f"Updated {len(documents)} documents with new catalog ID.")

    except Exception as e:
        logger.error(f"An error occurred during migration: {e}")
        session.rollback()
        raise
    finally:
        session.close()

    logger.info("Migration completed successfully.")


def downgrade():
    pass
