from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from mockmesh.storage.manager import StorageManager


def _collection_key(project: str, database: str, collection: str) -> str:
    return f"fs_{project}_{database}_{collection}"


def _doc_name(project: str, database: str, collection: str, doc_id: str) -> str:
    return f"projects/{project}/databases/{database}/documents/{collection}/{doc_id}"


def handle_firestore(
    sm: "StorageManager",
    project: str,
    database: str,
    collection: str,
    doc_id: str,
    operation: str,
    body: dict,
    tpl: dict,
) -> Optional[dict]:
    col_key = _collection_key(project, database, collection)

    if operation == "createDocument":
        doc_id = doc_id or str(uuid.uuid4())
        fields = body.get("fields", {})
        sm.nosql.put(col_key, doc_id, {"id": doc_id, "fields": fields})
        resp = dict(tpl.get("createDocument", {}))
        resp["name"] = _doc_name(project, database, collection, doc_id)
        resp["fields"] = fields
        return resp

    if operation == "getDocument":
        item = sm.nosql.get(col_key, doc_id)
        resp = dict(tpl.get("getDocument", {}))
        resp["name"] = _doc_name(project, database, collection, doc_id)
        if item:
            resp["fields"] = item.get("fields", {})
        return resp

    if operation == "updateDocument":
        existing = sm.nosql.get(col_key, doc_id) or {}
        fields   = {**existing.get("fields", {}), **body.get("fields", {})}
        sm.nosql.put(col_key, doc_id, {"id": doc_id, "fields": fields})
        resp = dict(tpl.get("updateDocument", {}))
        resp["name"]   = _doc_name(project, database, collection, doc_id)
        resp["fields"] = fields
        return resp

    if operation == "deleteDocument":
        sm.nosql.delete(col_key, doc_id)
        return {}

    if operation == "listDocuments":
        records = sm.nosql.list(col_key)
        resp    = dict(tpl.get("listDocuments", {"documents": []}))
        resp["documents"] = [
            {
                "name":   _doc_name(project, database, collection, r["key"]),
                "fields": r["value"].get("fields", {}),
            }
            for r in records
        ]
        return resp

    if operation == "runQuery":
        records = sm.nosql.list(col_key)
        resp_list = []
        for r in records:
            resp_list.append({
                "document": {
                    "name":   _doc_name(project, database, collection, r["key"]),
                    "fields": r["value"].get("fields", {}),
                }
            })
        return resp_list if resp_list else [{}]

    if operation == "batchWrite":
        resp = dict(tpl.get("batchWrite", {"writeResults": [], "status": []}))
        writes = body.get("writes", [])
        for write in writes:
            if "update" in write:
                doc  = write["update"]
                name = doc.get("name", "")
                did  = name.split("/")[-1] if name else str(uuid.uuid4())
                sm.nosql.put(col_key, did, {"id": did, "fields": doc.get("fields", {})})
                resp["writeResults"].append({"updateTime": "2024-01-01T00:00:00Z"})
            elif "delete" in write:
                did = write["delete"].split("/")[-1]
                sm.nosql.delete(col_key, did)
                resp["writeResults"].append({})
        return resp

    return None
