from pymongo import AsyncMongoClient
from pymongo.asynchronous.database import AsyncDatabase

from .config import settings

_client: AsyncMongoClient | None = None


async def connect() -> None:
    """Open the pool and fail fast if Mongo is unreachable."""
    global _client
    _client = AsyncMongoClient(settings.mongo_uri, serverSelectionTimeoutMS=5000)
    await _client.admin.command("ping")
    await ensure_indexes()


async def ensure_indexes() -> None:
    """Indexes for collections this API owns.

    Deliberately limited to `preprintsource`, which this API created. The
    pre-existing `users` collection is left alone -- adding a unique index there
    could fail on live data, so it is documented in the README instead.
    """
    db = get_db()
    await db["preprintsource"].create_index("title")
    await db["preprintsource"].create_index([("type", 1), ("frequency", 1)])
    await db["preprintsource"].create_index("created_on")
    # Drives the "still free to assign" lookup and each user's own list.
    await db["preprintsource"].create_index("assigned_to")

    harvested = db["collected_data_preprints_v2"]
    # title + email is the identity: the same author on the same paper is one
    # record. Enforced by the database, not just by the insert path.
    await harvested.create_index([("title", 1), ("email", 1)], unique=True)
    await harvested.create_index("email")
    await harvested.create_index("collected_on")
    await harvested.create_index("published_on")
    # Set by the block-list filters; keeps the per-reason / remaining counts fast.
    await harvested.create_index("filter_reason")

    # preprints.org has its OWN harvested collection (never mixes with arXiv).
    pp = db["collected_data_preprints_org"]
    await pp.create_index([("title", 1), ("email", 1)], unique=True)
    await pp.create_index("published_on")
    await pp.create_index("collected_on")
    await db["preprints_auto_runs"].create_index("date", unique=True)

    # bioRxiv has its OWN harvested collection too (never mixes with arXiv).
    bx = db["collected_data_biorxiv"]
    await bx.create_index([("title", 1), ("email", 1)], unique=True)
    await bx.create_index("published_on")
    await bx.create_index("collected_on")
    await bx.create_index("discipline")
    await bx.create_index("filter_reason")
    await db["biorxiv_auto_runs"].create_index("date", unique=True)

    # medRxiv has its OWN harvested collection too (never mixes with the others).
    mx = db["collected_data_medrxiv"]
    await mx.create_index([("title", 1), ("email", 1)], unique=True)
    await mx.create_index("published_on")
    await mx.create_index("collected_on")
    await mx.create_index("discipline")
    await mx.create_index("filter_reason")
    await db["medrxiv_auto_runs"].create_index("date", unique=True)

    # Processed output: drives the day-wise "saved" flag and skip-if-saved check.
    await db["collected_data_de"].create_index("published_on")

    # Publisher-wise campaigns and their copied records.
    await db["cd_campaigns"].create_index(
        [("user_id", 1), ("publisher_id", 1), ("primary_topic_id", 1), ("published_on", 1)],
        unique=True,
    )
    await db["cd_campaign_data"].create_index([("campaign_id", 1), ("de_id", 1)], unique=True)

    auto_runs = db["arxiv_auto_runs"]
    # One record per collected day.
    await auto_runs.create_index("date", unique=True)
    await auto_runs.create_index("status")
    await harvested.create_index("collected_by")
    await harvested.create_index("country")
    # Websites are the natural key; collation strength 2 makes it case-insensitive.
    await db["preprintsource"].create_index(
        "website", unique=True, collation={"locale": "en", "strength": 2}
    )



async def disconnect() -> None:
    global _client
    if _client is not None:
        await _client.close()
        _client = None


def get_db() -> AsyncDatabase:
    if _client is None:
        raise RuntimeError("Mongo client is not connected")
    return _client[settings.mongo_db]


def get_users():
    return get_db()["users"]