from pymongo.mongo_client import MongoClient
from pymongo.operations import SearchIndexModel
import pymongo

from typing import List
from pydantic import BaseModel, ValidationError
from transformers import AutoTokenizer, AutoModel
from tqdm import tqdm
import json
import time
import argparse
import logging
import os
from pathlib import Path
import torch
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())
# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# Check for local model first, then fall back to remote
model_name = "facebook/contriever"
# local_model_path = os.getenv("HF_MODEL_PATH") or str(
#     Path(__file__).parent.parent.parent.parent / "models" / "facebook--contriever"
# )

# if os.path.exists(local_model_path) and os.path.isdir(local_model_path):
#     model_path = local_model_path
# else:
model_path = model_name

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, use_safetensors=True).to(device)
# model = AutoModel.from_pretrained(model_path).to(device)


class EntityType(BaseModel):
    _id: int
    entity_type_id: str
    label: str
    parent_type_ids: List[str]
    valid_subject_property_ids: List[str]
    valid_object_property_ids: List[str]


class Property(BaseModel):
    _id: int
    property_id: str
    label: str
    valid_subject_type_ids: List[str]
    valid_object_type_ids: List[str]


class EntityTypeAlias(BaseModel):
    _id: int
    entity_type_id: str
    alias_label: str
    alias_text_embedding: List[float]


class PropertyAlias(BaseModel):
    _id: int
    relation_id: str
    alias_label: str
    alias_text_embedding: List[float]


def get_embedding(text):
    def mean_pooling(token_embeddings, mask):
        token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.0)
        sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]
        return sentence_embeddings

    if not text or not isinstance(text, str):
        return None

    try:
        inputs = tokenizer([text], padding=True, truncation=True, return_tensors="pt")
        outputs = model(**inputs.to("cuda"))
        embeddings = mean_pooling(outputs[0], inputs["attention_mask"])
        return embeddings.detach().cpu().tolist()[0]

    except Exception as e:
        logger.error(f"Error in get_embedding: {e}")
        return None


def get_mongo_client(mongo_uri):
    client = MongoClient(mongo_uri)
    logger.info("Connection to MongoDB successful")
    return client


def populate_entity_types(
    ENTITY_2_LABEL,
    ENTITY_2_HIERARCHY,
    SUBJ_2_PROP_CONSTRAINTS,
    OBJ_2_PROP_CONSTRAINTS,
    db,
    collection_name="entity_types",
):
    logger.info(f"Starting to populate {collection_name} collection")
    entity_metadata_list = []

    for i, entity in enumerate(ENTITY_2_LABEL.keys()):
        label = ENTITY_2_LABEL[entity]
        parents = ENTITY_2_HIERARCHY[entity]

        valid_subject_property_ids = (
            SUBJ_2_PROP_CONSTRAINTS[entity] if entity in SUBJ_2_PROP_CONSTRAINTS else []
        )
        valid_object_property_ids = (
            OBJ_2_PROP_CONSTRAINTS[entity] if entity in OBJ_2_PROP_CONSTRAINTS else []
        )

        entity_metadata_list.append(
            {
                "_id": i,
                "entity_type_id": entity,
                "label": label,
                "parent_type_ids": parents,
                "valid_subject_property_ids": valid_subject_property_ids,
                "valid_object_property_ids": valid_object_property_ids,
            }
        )

    entity_metadata_list.append(
        {
            "_id": i + 1,
            "entity_type_id": "ANY",
            "label": "ANY",
            "parent_type_ids": [],
            "valid_subject_property_ids": SUBJ_2_PROP_CONSTRAINTS["<ANY SUBJECT>"],
            "valid_object_property_ids": OBJ_2_PROP_CONSTRAINTS["<ANY OBJECT>"],
        }
    )

    try:
        records = [EntityType(**record).model_dump() for record in entity_metadata_list]
    except ValidationError as e:
        logger.error(f"Validation error while populating {collection_name}: {e}")

    collection = db.get_collection(collection_name)
    collection.insert_many(records)
    logger.info(f"Successfully populated {collection_name} with {len(records)} records")


def populate_entity_type_aliases(
    ENTITY_2_LABEL, ENTITY_2_ALIASES, db, collection_name="entity_type_aliases"
):
    logger.info(f"Starting to populate {collection_name} collection")
    entity_types_list = []
    id_count = 0

    for e, aliases in tqdm(ENTITY_2_ALIASES.items()):
        alias_embedding = get_embedding(ENTITY_2_LABEL[e])
        entity_types_list.append(
            {
                "_id": id_count,
                "entity_type_id": e,
                "alias_label": ENTITY_2_LABEL[e],
                "alias_text_embedding": alias_embedding,
            }
        )
        id_count += 1

        for alias in aliases:
            alias_embedding = get_embedding(alias)
            entity_types_list.append(
                {
                    "_id": id_count,
                    "entity_type_id": e,
                    "alias_label": alias,
                    "alias_text_embedding": alias_embedding,
                }
            )
            id_count += 1
    try:
        records = [
            EntityTypeAlias(**record).model_dump() for record in entity_types_list
        ]
    except ValidationError as e:
        logger.error(f"Validation error while populating {collection_name}: {e}")

    collection = db.get_collection(collection_name)
    collection.insert_many(records)
    logger.info(f"Successfully populated {collection_name} with {len(records)} records")


def populate_properties(
    PROP_2_LABEL, PROP_2_CONSTRAINT, db, collection_name="properties"
):
    logger.info(f"Starting to populate {collection_name} collection")
    property_list = []

    for i, prop_id in enumerate(PROP_2_LABEL.keys()):
        property_list.append(
            {
                "_id": i,
                "property_id": prop_id,
                "label": PROP_2_LABEL[prop_id],
                "valid_subject_type_ids": PROP_2_CONSTRAINT[prop_id][
                    "Subject type constraint"
                ],
                "valid_object_type_ids": PROP_2_CONSTRAINT[prop_id][
                    "Value-type constraint"
                ],
            }
        )

    try:
        records = [Property(**record).model_dump() for record in property_list]
    except ValidationError as e:
        logger.error(f"Validation error while populating {collection_name}: {e}")

    collection = db.get_collection(collection_name)
    collection.insert_many(records)
    logger.info(f"Successfully populated {collection_name} with {len(records)} records")


def populate_property_aliases(
    PROP_2_LABEL, PROP_2_ALIASES, db, collection_name="property_aliases"
):
    logger.info(f"Starting to populate {collection_name} collection")
    relation_alias_id_pairs = []
    id_count = 0

    for r, aliases in tqdm(PROP_2_ALIASES.items()):
        alias_embedding = get_embedding(PROP_2_LABEL[r])
        relation_alias_id_pairs.append(
            {
                "_id": id_count,
                "relation_id": r,
                "alias_label": PROP_2_LABEL[r],
                "alias_text_embedding": alias_embedding,
            }
        )
        id_count += 1

        for alias in aliases:
            alias_embedding = get_embedding(alias)
            relation_alias_id_pairs.append(
                {
                    "_id": id_count,
                    "relation_id": r,
                    "alias_label": alias,
                    "alias_text_embedding": alias_embedding,
                }
            )
            id_count += 1
    try:
        records = [
            PropertyAlias(**record).model_dump() for record in relation_alias_id_pairs
        ]
    except ValidationError as e:
        logger.error(f"Validation error while populating {collection_name}: {e}")

    collection = db.get_collection(collection_name)
    collection.insert_many(records)
    logger.info(f"Successfully populated {collection_name} with {len(records)} records")


def create_search_index_for_entity_types(
    db,
    collection_name="entity_type_aliases",
    embedding_field_name="alias_text_embedding",
    index_name="entity_type_aliases",
):
    logger.info(f"Starting to create index {index_name} for {collection_name}")
    collection = db.get_collection(collection_name)
    vector_search_index_model = SearchIndexModel(
        definition={
            "mappings": {
                "dynamic": True,
                "fields": {
                    embedding_field_name: {
                        "dimensions": 768,
                        "similarity": "cosine",
                        "type": "knnVector",
                    }
                },
            }
        },
        name=index_name,
    )

    try:
        result = collection.create_search_index(model=vector_search_index_model)
        logger.info("Creating index...")
        time.sleep(20)
        logger.info(f"New index {index_name} created successfully: {result}")
    except Exception as e:
        logger.error(f"Error creating new vector search index {index_name}: {str(e)}")


def create_search_index_for_properties(
    db,
    collection_name="property_aliases",
    embedding_field_name="alias_text_embedding",
    prop_id_field_name="relation_id",
    index_name="property_aliases_ids",
):
    logger.info(f"Starting to create index {index_name} for {collection_name}")
    collection = db.get_collection(collection_name)
    vector_search_index_model = SearchIndexModel(
        definition={
            "mappings": {
                "dynamic": True,
                "fields": {
                    embedding_field_name: {
                        "dimensions": 768,
                        "similarity": "cosine",
                        "type": "knnVector",
                    },
                    prop_id_field_name: {"type": "token"},
                },
            }
        },
        name=index_name,
    )

    try:
        result = collection.create_search_index(model=vector_search_index_model)
        logger.info("Creating index...")
        time.sleep(20)
        logger.info(f"New index {index_name} created successfully: {result}")
    except Exception as e:
        logger.error(f"Error creating new vector search index {index_name}: {str(e)}")


def create_indexes(db):
    logger.info("Creating indexes for entity_types collection...")
    db.entity_types.create_index([("entity_type_id", 1)])
    db.entity_types.create_index([("label", 1)])

    logger.info("Creating indexes for entity_type_aliases collection...")
    db.entity_type_aliases.create_index([("entity_type_id", 1)])
    db.entity_type_aliases.create_index([("alias_label", 1)])

    logger.info("Creating indexes for properties collection...")
    db.properties.create_index([("property_id", 1)])

    # logger.info("Creating indexes for property_aliases collection...")
    # db.property_aliases.create_index("relation_id")

    logger.info("Creating indexes for entity_aliases collection...")
    db.entity_aliases.create_index([("entity_type", 1), ("sample_id", 1)])
    db.entity_aliases.create_index([("label", 1)])

    db.create_collection("triplets")
    logger.info("Creating indexes for triplets collection...")
    db.triplets.create_index([("sample_id", 1)])

    logger.info("All indexes created successfully")


def create_wikidata_ontology_database(
    mongo_uri: str = "mongodb://localhost:27018/?directConnection=true",
    database: str = "wikidata_ontology",
    mappings_dir: str = None,
    entity_types_collection: str = "entity_types",
    entity_type_aliases_collection: str = "entity_type_aliases",
    properties_collection: str = "properties",
    property_aliases_collection: str = "property_aliases",
    entity_types_index: str = "entity_type_aliases",
    property_aliases_index: str = "property_aliases",
    drop_collections: bool = True,
):
    """
    Populate MongoDB with Wikidata ontology data.

    Args:
        mongo_uri: MongoDB connection URI
        database: MongoDB database name
        mappings_dir: Directory containing ontology mapping files. If None, uses default path.
        entity_types_collection: Collection name for entity types
        entity_type_aliases_collection: Collection name for entity type aliases
        properties_collection: Collection name for properties
        property_aliases_collection: Collection name for property aliases
        entity_types_index: Index name for entity types
        property_aliases_index: Index name for property aliases
        drop_collections: Whether to drop existing collections before creating new ones

    Returns:
        Database object
    """

    # Default mappings directory
    if mappings_dir is None:
        # Try to find the mappings directory relative to this file
        current_file = Path(__file__).parent
        mappings_dir = str(current_file / "utils" / "ontology_mappings" / "")
        if not os.path.exists(mappings_dir):
            # Fallback to relative path
            mappings_dir = "utils/ontology_mappings/"

    logger.info("Starting database population process")
    logger.info(f"Using database: {database}")
    logger.info(f"Loading mapping files from: {mappings_dir}")

    # Load mapping files
    with open(os.path.join(mappings_dir, "subj_constraint2prop.json"), "r") as f:
        subj2prop_constraints = json.load(f)

    with open(os.path.join(mappings_dir, "obj_constraint2prop.json"), "r") as f:
        obj2prop_constraints = json.load(f)

    with open(os.path.join(mappings_dir, "entity_type2label.json"), "r") as f:
        ENTITY_2_LABEL = json.load(f)

    with open(os.path.join(mappings_dir, "entity_type2hierarchy.json"), "r") as f:
        ENTITY_2_HIERARCHY = json.load(f)

    with open(os.path.join(mappings_dir, "entity_type2aliases.json"), "r") as f:
        ENTITY_2_ALIASES = json.load(f)

    with open(os.path.join(mappings_dir, "prop2constraints.json"), "r") as f:
        PROP_2_CONSTRAINT = json.load(f)

    with open(os.path.join(mappings_dir, "prop2label.json"), "r") as f:
        PROP_2_LABEL = json.load(f)

    with open(os.path.join(mappings_dir, "prop2aliases.json"), "r") as f:
        PROP_2_ALIASES = json.load(f)

    logger.info("Successfully loaded all mapping files")

    # Connect to MongoDB
    mongo_client = get_mongo_client(mongo_uri)
    db = mongo_client.get_database(database)

    # Drop all existing collections
    if drop_collections:
        logger.info("Dropping existing collections...")
        for collection_name in db.list_collection_names():
            logger.info(f"Dropping collection: {collection_name}")
            db.drop_collection(collection_name)
        logger.info("Successfully dropped all existing collections")

    # Populate collections
    populate_entity_types(
        ENTITY_2_LABEL,
        ENTITY_2_HIERARCHY,
        subj2prop_constraints,
        obj2prop_constraints,
        db,
        collection_name=entity_types_collection,
    )

    populate_entity_type_aliases(
        ENTITY_2_LABEL,
        ENTITY_2_ALIASES,
        db,
        collection_name=entity_type_aliases_collection,
    )

    populate_properties(
        PROP_2_LABEL, PROP_2_CONSTRAINT, db, collection_name=properties_collection
    )

    populate_property_aliases(
        PROP_2_LABEL,
        PROP_2_ALIASES,
        db,
        collection_name=property_aliases_collection,
    )

    # Create search indexes
    create_search_index_for_entity_types(
        db,
        collection_name=entity_type_aliases_collection,
        index_name=entity_types_index,
    )
    create_search_index_for_properties(
        db,
        collection_name=property_aliases_collection,
        index_name=property_aliases_index,
    )

    # Create indexes
    create_indexes(db)
    logger.info("Database population process completed")

    return db


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Populate MongoDB with Wikidata ontology data"
    )

    parser.add_argument(
        "--mappings_dir",
        type=str,
        default="utils/ontology_mappings/",
        help="Directory containing ontology mapping files",
    )
    parser.add_argument(
        "--mongo_uri",
        type=str,
        default="mongodb://localhost:27018/?directConnection=true",
        help="MongoDB connection URI",
    )
    parser.add_argument(
        "--database",
        type=str,
        default="wikidata_ontology",
        help="MongoDB database name",
    )

    # Collection names
    parser.add_argument(
        "--entity_types_collection",
        type=str,
        default="entity_types",
        help="Collection name for entity types",
    )
    parser.add_argument(
        "--entity_type_aliases_collection",
        type=str,
        default="entity_type_aliases",
        help="Collection name for entity type aliases",
    )
    parser.add_argument(
        "--properties_collection",
        type=str,
        default="properties",
        help="Collection name for properties",
    )
    parser.add_argument(
        "--property_aliases_collection",
        type=str,
        default="property_aliases",
        help="Collection name for property aliases",
    )

    # Index names
    parser.add_argument(
        "--entity_types_index",
        type=str,
        default="entity_type_aliases",
        help="Index name for entity types",
    )
    parser.add_argument(
        "--property_aliases_index",
        type=str,
        default="property_aliases",
        help="Index name for property aliases",
    )

    args = parser.parse_args()
    create_wikidata_ontology_database(
        mongo_uri=args.mongo_uri,
        database=args.database,
        mappings_dir=args.mappings_dir,
        entity_types_collection=args.entity_types_collection,
        entity_type_aliases_collection=args.entity_type_aliases_collection,
        properties_collection=args.properties_collection,
        property_aliases_collection=args.property_aliases_collection,
        entity_types_index=args.entity_types_index,
        property_aliases_index=args.property_aliases_index,
    )
