import json
import logging
import os
import sys
from typing import Any, Dict, List, Tuple, Union

from dotenv import load_dotenv
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    Function,
    FunctionType,
    connections,
    utility,
)

logger = logging.getLogger(__name__)


def load_env() -> None:
    """Загрузка переменных окружения из .env файла."""
    load_dotenv()
    logger.debug('Переменные окружения загружены из .env файла')


def connect_to_milvus() -> None:
    """Подключение к Milvus."""
    milvus_host = os.getenv('MILVUS_HOST')
    milvus_port = os.getenv('MILVUS_PORT')
    database_name = os.getenv('DATABASE_NAME')

    if not all([milvus_host, milvus_port, database_name]):
        logger.error('Ошибка: Не заданы переменные окружения MILVUS_HOST, MILVUS_PORT или DATABASE_NAME')
        sys.exit(1)

    try:
        connections.connect(host=milvus_host, port=milvus_port, db_name=database_name)
        logger.debug(f'Детали подключения: {connections}')
    except Exception as error:
        logger.error(f'Ошибка подключения к Milvus: {error}')
        sys.exit(1)


def make_safe_db_name(project_id: str) -> str:
    """
    Создает безопасное имя базы данных из project_id
    PostgreSQL требует, чтобы имена БД содержали только буквы, цифры и подчеркивания
    """
    safe_id = project_id.replace('-', '_')
    #safe_id = ''.join(c if c.isalnum() or c == '_' else '_' for c in safe_id)

    return f"xml_{safe_id}"


def create_collection(project_id: str) -> Collection:
    """Создание коллекции в Milvus для хранения информации о таблицах БД."""
    collection_name = make_safe_db_name(project_id)
    if utility.has_collection(collection_name):
        logger.info(f'Коллекция {collection_name} уже существует.')
        return Collection(name=collection_name)

    fields = [
        FieldSchema(name='table_name', dtype=DataType.VARCHAR, max_length=255, is_primary=True),
        FieldSchema(name='table_info_raw', dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(name="file_id", dtype=DataType.VARCHAR, max_length=256),
        FieldSchema(name='file_title', dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(
            name='bm25_text',
            dtype=DataType.VARCHAR,
            max_length=65535,
            enable_analyzer=True,
            analyzer_params={'tokenizer': 'whitespace'},
        ),
        FieldSchema(name='embedding_bm25', dtype=DataType.SPARSE_FLOAT_VECTOR),
        FieldSchema(name='num_fields', dtype=DataType.INT64),
        FieldSchema(name='table_info', dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(name='related_tables', dtype=DataType.VARCHAR, max_length=65535),
    ]

    schema = CollectionSchema(fields)

    bm25_function = Function(
        name='text_bm25_emb',
        input_field_names=['bm25_text'],
        output_field_names=['embedding_bm25'],
        function_type=FunctionType.BM25,
    )
    schema.add_function(bm25_function)

    try:
        collection = Collection(name=collection_name, schema=schema)

        index_params_bm25 = {
            'index_type': 'SPARSE_INVERTED_INDEX',
            'metric_type': 'BM25',
            'params': {'inverted_index_algo': 'DAAT_MAXSCORE', 'bm25_k1': 1.2, 'bm25_b': 0.3},
        }
        collection.create_index(field_name='embedding_bm25', index_params=index_params_bm25)

        logger.info(f'Коллекция {collection_name} успешно создана с индексом.')
        return collection
    except Exception as error:
        logger.error(f'Ошибка при создании коллекции: {error}')
        sys.exit(1)


def upsert_tables_data(collection: Collection, chunks: List[Dict[str, Any]]) -> Tuple[int, int, int]:
    """
    Вставляет или обновляет данные в коллекции Milvus по логике:
    1. Если таблицы с таким table_name нет, добавить новую запись
    2. Если таблица уже существует:
       - Если у новой версии num_fields больше, заменить старую запись
         с объединением связанных таблиц (related_tables)
       - Если у новой версии num_fields меньше или равно, оставить старую запись

    Args:
        collection: Коллекция Milvus
        chunks: Список словарей с данными о таблицах

    Returns:
        tuple: (количество вставленных записей, количество обновленных записей, количество пропущенных записей)
    """
    if not chunks:
        logger.warning('Пустой список для вставки/обновления')
        return 0, 0, 0

    table_names = [chunk['table_name'] for chunk in chunks]

    existing_records = {}
    existing_relations = {}  # Добавим словарь для хранения существующих связей

    try:
        collection.load()
        logger.debug('Коллекция успешно загружена в память')

        if table_names:
            # Формируем корректное выражение для запроса
            if len(table_names) == 1:
                # Для одной таблицы используем равенство
                expr = f"table_name == '{table_names[0]}'"
            else:
                # Для нескольких таблиц используем IN
                names_str = "['" + "','".join(table_names) + "']"
                expr = f'table_name in {names_str}'

            logger.debug(f'Выражение для запроса: {expr}')
            # Добавляем related_tables в output_fields для получения существующих связей
            results = collection.query(expr=expr, output_fields=['table_name', 'num_fields', 'related_tables'])

            # Создаем словарь существующих записей для быстрого доступа
            for record in results:
                existing_records[record['table_name']] = record['num_fields']
                # Сохраняем существующие связи
                if 'related_tables' in record:
                    try:
                        existing_relations[record['table_name']] = json.loads(record['related_tables'])
                    except (json.JSONDecodeError, TypeError):
                        logger.warning(f"Не удалось декодировать related_tables для {record['table_name']}")
                        existing_relations[record['table_name']] = []

            logger.debug(f'Найдено {len(existing_records)} существующих записей в коллекции')
        else:
            logger.debug('Список таблиц пуст, пропускаем запрос')
    except Exception as error:
        logger.error(f'Ошибка при поиске существующих записей: {error}')
        logger.error(f'Детали ошибки: {str(error)}')
        # В случае ошибки считаем, что существующих записей нет
        pass

    # Разделяем данные на новые записи и кандидаты на обновление
    new_records = []
    records_to_update = []
    skipped_records = []

    for chunk in chunks:
        table_name = chunk['table_name']
        num_fields = chunk['num_fields']

        if table_name not in existing_records:
            # Таблица не существует - добавляем
            new_records.append(chunk)
            logger.debug(f"Таблица '{table_name}' будет добавлена как новая запись")
        else:
            # Таблица существует - проверяем num_fields
            existing_num_fields = existing_records[table_name]

            if num_fields >= existing_num_fields:
                # Новая версия содержит больше полей - обновляем

                # Объединяем related_tables
                if table_name in existing_relations:
                    try:
                        # Получаем текущие связи из chunk
                        current_relations = json.loads(chunk['related_tables'])
                        # Объединяем существующие и новые связи
                        merged_relations = list(set(existing_relations[table_name] + current_relations))
                        # Обновляем значение в chunk
                        chunk['related_tables'] = json.dumps(merged_relations, ensure_ascii=False)
                        logger.debug(
                            f"Объединены связи для таблицы '{table_name}': было {len(existing_relations[table_name])}, добавлено {len(current_relations)}, стало {len(merged_relations)}",
                        )
                    except (json.JSONDecodeError, TypeError) as error:
                        logger.warning(f'Ошибка при объединении related_tables для {table_name}: {error}')

                records_to_update.append(chunk)
                logger.debug(
                    f"Таблица '{table_name}' будет обновлена (новых полей: {num_fields - existing_num_fields})",
                )
            else:
                # Старая версия содержит не меньше полей - пропускаем
                skipped_records.append(table_name)
                logger.debug(
                    f"Таблица '{table_name}' пропущена (существующих полей: {existing_num_fields}, новых: {num_fields})",
                )

    inserted_count = 0
    updated_count = 0

    # Вставляем новые записи
    if new_records:
        try:
            insert_result = collection.insert(new_records)
            inserted_count = len(new_records)
            logger.info(f'Добавлено {inserted_count} новых записей')
        except Exception as error:
            logger.error(f'Ошибка при вставке новых записей: {error}')

    # Обновляем существующие записи, которые нужно заменить
    if records_to_update:
        try:
            # Сначала удаляем старые записи
            for record in records_to_update:
                expr = f"table_name == '{record['table_name']}'"
                collection.delete(expr)

            # Затем вставляем обновленные
            update_result = collection.insert(records_to_update)
            updated_count = len(records_to_update)
            logger.info(f'Обновлено {updated_count} существующих записей')
        except Exception as error:
            logger.error(f'Ошибка при обновлении записей: {error}')

    # Фиксируем изменения
    collection.flush()

    try:
        collection.release()
        logger.debug('Коллекция успешно освобождена из памяти')
    except Exception as error:
        logger.warning(f'Не удалось освободить коллекцию: {str(error)}')

    return inserted_count, updated_count, len(skipped_records)


def parse_schema_to_api_format(schema_dict: Dict[str, Dict[str, str]]) -> List[Dict[str, Union[str, List[str]]]]:
    """Преобразует словарь схемы базы данных в формат для API векторизатора.

    Args:
        schema_dict: Словарь в формате:
            {'table_name': {'PK column_name': 'data_type NOT NULL', 'column_name': 'data_type'}, ...}

    Returns:
        Список словарей в формате:
            [{"table_name": "table_name", "fields": ["column_name1", "column_name2", ...]}, ...]
    """
    result = []

    for table_name, columns in schema_dict.items():
        fields = []
        for column_name in columns.keys():
            clean_column_name = column_name.replace('PK ', '').replace('FK ', '')
            fields.append(clean_column_name)

        result.append({'table_name': table_name, 'fields': fields})

    logger.debug(f'Схема преобразована в API формат для {len(result)} таблиц')
    return result


def parsed_final(table_name: str, table_info_raw: Dict[str, str]) -> str:
    parsed_dict = parse_schema_to_api_format({table_name: table_info_raw})

    texts = []
    for schema_entry in parsed_dict:
        repeated_table = f"{schema_entry['table_name']} {schema_entry['table_name']} {schema_entry['table_name']}"
        fields_text = ' '.join(schema_entry['fields'])
        combined_text = f'{repeated_table} {fields_text}'
        texts.append(combined_text)

    return texts[0]


def json_to_text_schema(json_schema: Dict[str, Dict[str, Any]]) -> List[str]:
    """Преобразует JSON схему в текстовое представление."""
    text_schema = []

    for table_name, table_fields in json_schema.items():
        # Извлекаем первичный ключ
        primary_key = next((key.split()[1] for key in table_fields.keys() if key.startswith('PK ')), 'N/A')
        primary_key_type = table_fields.get(f'PK {primary_key}', 'N/A')

        # Подготавливаем список полей
        fields = [
            f'* {field}: {field_type}' for field, field_type in table_fields.items() if not field.startswith('PK ')
        ]

        # Формируем описание таблицы в требуемом формате
        fields_str = '\n'.join(fields)
        table_description = (
            f'Table: {table_name}\n- Primary Key: {primary_key} ({primary_key_type})\n- Fields:\n{fields_str}'
        )

        text_schema.append(table_description)

    logger.debug(f'Преобразовано {len(text_schema)} таблиц в текстовое представление')
    return text_schema


def prepare_chunks(
    schema_dict: Dict[str, Dict[str, str]],
    text_schema: List[str],
    related_tables: Dict[str, List[str]],
    file_id: str,
    file_title: str
) -> List[Dict[str, Any]]:
    """Подготавливает данные для вставки в Milvus.

    Args:
        schema_dict: Словарь схемы базы данных
        text_schema: Текстовое представление схемы
        related_tables: Словарь связанных таблиц
    """
    chunks = []

    for idx, (table_name, table_info) in enumerate(schema_dict.items()):
        related = related_tables.get(table_name, [])

        related_json = json.dumps(related, ensure_ascii=False)
        num_fields = len(table_info.keys())

        chunk = {
            'table_name': table_name,
            'table_info_raw': str(table_info),
            'file_id': file_id,
            'file_title': file_title,
            'bm25_text': parsed_final(table_name, table_info),
            'num_fields': num_fields,
            'table_info': text_schema[idx],
            'related_tables': related_json,
        }
        chunks.append(chunk)

    return chunks
