from __future__ import annotations

from typing import Dict, Optional

from autotouch_cli.exceptions import AutotouchInputError

try:
    from pymongo import MongoClient  # type: ignore
except Exception:  # pragma: no cover
    MongoClient = None


def mongo_client(uri: Optional[str], db_name: Optional[str]):
    if MongoClient is None:
        raise AutotouchInputError("pymongo not installed. Install requirements and retry.")
    import os

    resolved_uri = uri or os.environ.get("MONGODB_URI")
    if not resolved_uri:
        raise AutotouchInputError("MONGODB_URI not set. Export MONGODB_URI or pass --mongo-uri")
    resolved_db = db_name or os.environ.get("MONGODB_DB_NAME") or "autotouch"
    client = MongoClient(resolved_uri)
    return client[resolved_db]


def status_counts(db, table_id: str, column_id: str) -> Dict[str, int]:
    row_ids = [str(r["_id"]) for r in db["rows"].find({"table_id": table_id}, {"_id": 1})]
    if not row_ids:
        return {"pending": 0, "done": 0, "error": 0}
    pipeline = [
        {"$match": {"column_id": column_id, "row_id": {"$in": row_ids}}},
        {"$group": {"_id": {"$ifNull": ["$status", "missing"]}, "count": {"$sum": 1}}},
    ]
    results = list(db["cells"].aggregate(pipeline))
    counts = {doc.get("_id"): int(doc.get("count") or 0) for doc in results}
    for key in ("pending", "done", "error"):
        counts.setdefault(key, 0)
    return {key: counts[key] for key in ("pending", "done", "error")}
