from __future__ import annotations

import uuid
from typing import TYPE_CHECKING, Optional

from mockmesh.interceptors.azure.utils import _enc, _get_body, _safe_json
from mockmesh.interceptors.responses import FakeAzureResponse

if TYPE_CHECKING:
    from mockmesh.storage.manager import StorageManager


def handle_cosmos(
    sm: "StorageManager", parsed: dict, request: object,
    method: str, url: str, tpl: dict,
) -> Optional[FakeAzureResponse]:
    collection = parsed["container"]
    doc_id     = parsed["path"]
    raw_body   = _get_body(request)
    body_data  = _safe_json(raw_body)

    if method == "GET" and not collection:
        discovery = dict(tpl.get("GET") or tpl.get("discovery_fallback") or {})
        actual_ep = url.split("/dbs")[0].rstrip("/") + "/"
        for loc in discovery.get("writableLocations", []):
            loc["databaseAccountEndpoint"] = actual_ep
        for loc in discovery.get("readableLocations", []):
            loc["databaseAccountEndpoint"] = actual_ep
        return FakeAzureResponse(200, {"Content-Type": "application/json"}, _enc(discovery))

    req_headers: dict = {}
    try:
        req_headers = dict(getattr(request, "headers", {}) or {})
    except Exception:
        pass
    is_query = (
        "query" in body_data
        or str(req_headers.get("x-ms-documentdb-isquery", "")).lower() == "true"
        or str(req_headers.get("x-ms-documentdb-query-enable-scan", "")).lower() == "true"
    )
    cosmos_col = f"cosmos--{collection}" if collection else ""

    if method == "POST" and is_query:
        items    = sm.nosql.list(cosmos_col)
        list_tpl = dict(tpl.get("POST", {"Documents": [], "_count": 0, "_rid": "mock"}))
        list_tpl["Documents"] = [r["value"] for r in items]
        list_tpl["_count"]    = len(items)
        return FakeAzureResponse(200, {"Content-Type": "application/json"}, _enc(list_tpl))

    if method in ("PUT", "POST") and body_data:
        doc_id = doc_id or str(body_data.get("id", uuid.uuid4()))
        body_data["id"] = doc_id
        sm.cosmos_upsert(collection, body_data)
        etag = '"' + doc_id[:8] + '-mock"'
        return FakeAzureResponse(tpl.get("UPSERT", {}).get("status", 201), {"ETag": etag}, _enc(body_data))

    if method == "GET" and doc_id:
        item = sm.cosmos_get(collection, doc_id)
        if item:
            etag = '"' + doc_id[:8] + '-mock"'
            return FakeAzureResponse(tpl.get("GET_doc", {}).get("status", 200), {"ETag": etag}, _enc(item))

    if method == "GET" and not doc_id:
        items    = sm.nosql.list(cosmos_col)
        list_tpl = dict(tpl.get("POST", {"Documents": [], "_count": 0, "_rid": "mock"}))
        list_tpl["Documents"] = [r["value"] for r in items]
        list_tpl["_count"]    = len(items)
        return FakeAzureResponse(200, {"Content-Type": "application/json"}, _enc(list_tpl))

    if method == "DELETE" and doc_id:
        sm.nosql.delete(cosmos_col, doc_id)
        return FakeAzureResponse(tpl.get("DELETE", {}).get("status", 204), {}, b"")

    return None
