import os
import logging
from fastapi import HTTPException
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    connections,
    utility,
)

logger = logging.getLogger(__name__)



def connect_to_milvus() -> bool:
    """Подключение к 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')
        return False

    try:
        connections.connect(host=milvus_host, port=milvus_port, db_name=database_name)
        logger.info(f'Подключение к Milvus установлено: {milvus_host}:{milvus_port}/{database_name}')
        return True
    except Exception as error:
        logger.error(f'Ошибка подключения к Milvus: {error}')
        return False


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"project_{safe_id}"

def create_collection_if_not_exists(project_id: str, vector_dim: int) -> 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)
    
    try:
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="page_id", dtype=DataType.VARCHAR, max_length=256),
            FieldSchema(name="link", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="page_title", dtype=DataType.VARCHAR, max_length=2048),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
            FieldSchema(name="chunk_index", dtype=DataType.INT64)
        ]
        
        schema = CollectionSchema(fields)
        collection = Collection(name=collection_name, schema=schema)
         
        index_params = {
            "metric_type": "COSINE",
        }
        collection.create_index("embedding", index_params)
        
        logger.info(f'Коллекция {collection_name} успешно создана с индексами.')
        return collection
        
    except Exception as error:
        logger.error(f'Ошибка при создании коллекции {collection_name}: {error}')
        raise HTTPException(status_code=500, detail=f"Не удалось создать коллекцию: {str(error)}")


def disconnect_from_milvus():
    """Отключение от Milvus"""
    try:
        connections.disconnect(alias='default')
        logger.debug('Отключение от Milvus выполнено')
    except Exception as e:
        logger.warning(f'Ошибка при отключении от Milvus: {str(e)}')
