import json
import typing
from typing import (
    AsyncIterator,
    Literal,
    Callable,
    Protocol,
    TypeVar,
    Generic,
    ClassVar,
    Any,
)
import warnings

import httpx

from . import structs


DOCT = TypeVar("DOCT")

# 3.12: type TypeIDType = ...
# Technically, anything JSONable is allowable, but only allowing atomic types
# makes a bunch of reasoning easier.
#: The type of class identifiers
TypeIDType = str | int | bool | None


class DocumentLoader(Protocol, Generic[DOCT]):
    """
    Protocol for things that go between JSON blobs and hydrated document
    instances.
    """

    def load_from_blob(self, blob: dict) -> DOCT:
        """
        Convert a JSON blob into a document object.
        """

    def dump_to_blob(self, doc: DOCT) -> dict:
        """
        Convert a document into a JSON blob.
        """

    def update_doc(self, doc: DOCT, **fields):
        """
        Update a document object in-place.
        """


class DocumentRegistry:
    """
    Handles de/serialization, manages migrations, etc.

    Default loader implementation.

    Do not use directly. You probably want one of the :ref:`integrations`.
    """

    TYPE_KEY = ""

    _docclasses: ClassVar[dict[TypeIDType, type]] = {}
    _migrations: ClassVar[list[tuple[TypeIDType, TypeIDType, Callable]]] = []

    def __init_sublcass__(cls):
        cls._docclasses = {}
        cls._migrations = []

    @classmethod
    def _get_class_from_name(cls, name: TypeIDType) -> type:
        return cls._docclasses[name]

    @classmethod
    def _get_name_from_class(cls, klass: type) -> TypeIDType:
        for name, kind in cls._docclasses.items():
            if issubclass(klass, kind):  # In case of decorator shenanigans
                return name
        else:
            raise ValueError(f"Couldn't find name for {klass}")

    @classmethod
    def document(cls, name: TypeIDType):
        """
        Register a class as a loadable couch document.

        Args:
            name: The type identifier to save to CouchDB

        .. note::

           The given name must be globally unique and must never change.
        """
        assert not isinstance(name, type)
        assert name not in cls._docclasses

        def _(klass: type):
            cls._docclasses[name] = klass
            return klass

        return _

    @classmethod
    def migration(cls, before: type, after: type):
        """
        Define a function that'll convert between documents.
        """
        # Normalize to the document classes previously registered
        bname = cls._get_name_from_class(before)
        aname = cls._get_name_from_class(after)
        # Enforce linearity
        assert not any(b == bname for b, _, _ in cls._migrations)

        def _(func: Callable):
            cls._migrations.append((bname, aname, func))
            return func

        return _

    def load_doc(self, cls: type, blob: dict):
        """
        Converts a JSON blob into a document.

        Override me.
        """
        raise NotImplementedError

    def dump_doc(self, doc) -> dict:
        """
        Convert a document into a JSON blob.

        Override me.
        """
        raise NotImplementedError

    def update_doc(self, doc, **fields):
        """
        Update a doc in-place
        """
        raise NotImplementedError

    def _migrate(self, bname, doc):
        while funcs := [f for b, _, f in self._migrations if b == bname]:
            (func,) = funcs
            doc = func(doc)
            bname = self._get_name_from_class(type(doc))
        return doc

    def load_from_blob(self, blob):
        type = blob.pop(self.TYPE_KEY)
        klass = self._get_class_from_name(type)
        doc = self.load_doc(klass, blob)
        doc = self._migrate(type, doc)
        return doc

    def dump_to_blob(self, doc):
        blob = self.dump_doc(doc)
        blob[self.TYPE_KEY] = self._get_name_from_class(type(doc))
        return blob


class Conflict(Exception):
    """
    There was a conflict when trying to perform the operation.
    """


class Missing(Exception):
    """
    Could not find the requested document.

    Note that this is a 404, not a tombstone.
    """


class Deleted(Exception):
    """
    Requested a deleted document.

    Note that this is a document with a tombstone, not a 404.
    """


class TooManyResults(Exception):
    """
    Requested a single document in a find operation and found more than one.

    Note this is based on number of results. Request was successful.
    """


class FindWarning(UserWarning):
    """
    Warnings reported by the CouchDB _find endpoint.
    """


class CouchSession:
    """
    A connection to CouchDB.

    You probably want to override like::

        class MySession(CouchSession):
            loader = MyRegistry

    """

    _client: httpx.AsyncClient
    _root: httpx.URL

    #: Class responsible for de/serializing data.
    loader: type[DocumentLoader]

    def __init__(self, client: httpx.AsyncClient, root: httpx.URL):
        self._client = client
        self._root = root

    @staticmethod
    def _fix_params(params):
        rv = {}
        for key, value in params.items():
            if value is None:
                continue
            elif isinstance(value, str):
                rv[key] = value
            else:
                rv[key] = json.dumps(value)
        return rv

    async def _request(self, method, *urlparts, **kwargs):
        url = self._root.join("/".join(urlparts))
        if "params" in kwargs:
            kwargs["params"] = self._fix_params(kwargs["params"])
        resp = await self._client.request(method, url, **kwargs)
        try:
            resp.raise_for_status()
        except httpx.HTTPStatusError as exc:
            exc.add_note(f"Body: {exc.response.text}")
            match exc.response.status_code:
                case 404:
                    raise Missing(f"Could not find {'/'.join(urlparts)}") from exc
                case 409:
                    raise Conflict(f"Conflict updating {'/'.join(urlparts)}") from exc
                case _:
                    raise
        return resp

    def __getitem__(self, key: str) -> "Database":
        """
        Gets a database.

        (Does not actually check if it exists.)
        """
        return Database(self, key)

    async def get_db(self, dbname: str) -> "Database":
        """
        Gets a database. Checks if it exists.

        See :http:head:`/{db}`
        """
        await self._request("HEAD", dbname)
        return Database(self, dbname)

    async def create_db(
        self,
        dbname: str,
        *,
        shards: int | None = None,
        replicas: int | None = None,
        partitioned: bool | None = None,
    ) -> "Database":
        """
        Create a database

        See :http:post:`/{db}`
        """
        try:
            await self._request(
                "PUT",
                dbname,
                params={
                    "q": shards,
                    "n": replicas,
                    "partitioned": partitioned,
                },
            )
        except httpx.HTTPStatusError as exc:
            match exc.response.status_code:
                case 412:
                    raise Conflict("Database already exists") from exc
                case _:
                    raise
        return Database(self, dbname)

    async def delete_db(self, dbname: str):
        """
        Delete a database.

        See :http:delete:`/{db}`
        """
        await self._request("DELETE", dbname)

    async def iter_dbs(self) -> AsyncIterator[str]:
        """
        List all databases

        See :http:get:`/_all_dbs`
        """
        resp = await self._request("GET", "_all_dbs")
        for dbname in resp.json():
            yield dbname

    # TODO: Database metadata


class Database:
    """
    An individual database.
    """

    def __init__(self, session, name):
        """
        :private:
        """
        self._session = session
        self._name = name

    def _blob2doc(self, blob, db, docid, etag=...):
        if docid is ...:
            docid = blob["_id"]
        if etag is ...:
            etag = f'"{blob["_rev"]}"'
        doc = self._session.loader().load_from_blob(blob)
        doc.__db = db
        doc.__docid = docid
        doc.__etag = etag
        return doc

    def _doc2blob(self, doc):
        blob = self._session.loader().dump_to_blob(doc)
        db = docid = etag = None
        try:
            db = doc.__db
            docid = doc.__docid
            etag = doc.__etag
        except AttributeError:
            pass
        return blob, db, docid, etag

    def _touch_doc(self, doc, *, db=None, etag=None, rev=None, id=None):
        fields = {}
        if etag is not None:
            doc.__etag = etag
        if db is not None:
            doc.__db = db
        if rev is not None:
            fields["_rev"] = rev
        if id is not None:
            doc.__docid = id
            fields["_id"] = id
        if fields:
            self._session.loader().update_doc(doc, **fields)

    async def get(
        self,
        docid: str,
        *,
        attachments: bool = False,
        conflicts: bool = False,
        deleted_conflicts: bool = False,
        latest: bool = False,
        local_seq: bool = False,
        meta: bool = False,
        open_revs: list[str] | Literal["all"] | None = None,
        rev: str | None = None,
        revs: bool = False,
        revs_info: bool = False,
    ):
        """
        Get a document

        See :http:get:`/{db}/{docid}`
        """
        resp = await self._session._request(
            "GET",
            self._name,
            docid,
            params={
                "attachments": attachments,
                "conflicts": conflicts,
                "deleted_conflicts": deleted_conflicts,
                "latest": latest,
                "local_seq": local_seq,
                "meta": meta,
                "open_revs": open_revs,
                "rev": rev,
                "revs": revs,
                "revs_info": revs_info,
            },
            headers={
                "Accept": "application/json",
            },
        )

        blob = resp.json()
        if blob.get("_deleted", False):  # TODO: Flag to override this
            raise Deleted("Document {self._name}/{docid} is marked as deleted")
        if "ETag" in resp.headers:
            etag = resp.headers["ETag"]
        else:
            # Conflicts mode
            etag = f'"{blob["_rev"]}"'

        # For some reason, requested fiels are omitted when empty
        if attachments:
            blob.setdefault("_attachments", {})
        if conflicts:
            blob.setdefault("_conflicts", [])
        if deleted_conflicts:
            blob.setdefault("_deleted_conflicts", [])
        if revs:
            blob.setdefault("_revisions", {})
        if revs_info or open_revs:
            blob.setdefault("_revisions", {})

        doc = self._blob2doc(blob, self._name, docid, etag)
        return doc

    # TODO: Attachments

    async def find_one(
        self, selector: typing.Mapping, use_index: str | list[str] | None = None
    ):
        """
        Get a single document based on ``selector``.

        See :http:post:`/{db}/_find`
        """
        json_body = {"selector": selector, "limit": 2}

        if use_index is not None:
            json_body |= {"use_index": use_index}

        resp = await self._session._request("POST", self._name, "_find", json=json_body)

        blob = resp.json()
        results = blob.get("docs")
        match len(results):
            case 0:
                raise Missing("No results found.")
            case 1:
                return self._blob2doc(results[0], self._name, ...)
            case _:
                raise TooManyResults("More than one result found.")

    async def find(
        self,
        selector: typing.Mapping,
        use_index: str | list[str] | None = None,
        pagesize: int | None = None,
    ) -> AsyncIterator:
        """
        Generate documents based on ``selector``.

        See :http:post:`/{db}/_find`
        """
        json_body: dict[str, Any] = {"selector": selector, "bookmark": None}

        if use_index is not None:
            json_body |= {"use_index": use_index}

        if pagesize is not None:
            json_body |= {"limit": pagesize}

        while True:
            resp = await self._session._request(
                "POST", self._name, "_find", json=json_body
            )
            payload = resp.json()
            if payload.get("warning", None):
                warnings.warn(payload["warning"], FindWarning)

            for doc in payload["docs"]:
                yield self._blob2doc(doc, self._name, ...)

            if not payload["docs"]:
                break

            json_body["bookmark"] = payload["bookmark"]

    async def attempt_put(
        self,
        doc,
        docid: str | None = None,
        *,
        batch: bool = False,
    ):
        """
        Update a document.

        db and docid only need to be given if it's a new document.

        See :http:put:`/{db}/{docid}`
        """
        blob, _db, _docid, etag = self._doc2blob(doc)
        assert _db is None or _db == self._name
        resp = await self._session._request(
            "PUT",
            self._name,
            _docid or docid,
            params={"batch": "ok"} if batch else {},
            headers={"If-Match": etag} if etag else {},
            json=blob,
        )
        payload = resp.json()
        assert payload["ok"]
        self._touch_doc(
            doc,
            db=self._name,
            id=payload["id"],
            etag=resp.headers["ETag"],
            rev=resp.json()["rev"],
        )

    async def attempt_delete(self, doc, *, batch: bool = False):
        """
        Delete a document

        See :http:delete:`/{db}/{docid}`
        """
        _, db, docid, etag = self._doc2blob(doc)
        assert db == self._name
        assert docid
        await self._session._request(
            "DELETE",
            db,
            docid,
            params={"batch": "ok"} if batch else {},
            headers={"If-Match": etag},
        )

    async def attempt_copy(self, src_doc, dst_doc, *, batch: bool = False):
        """
        Copy a document

        .. todo::

           Implement

        See :http:copy:`/{db}/{docid}`
        """
        # FIXME: Figure out signature

    async def mutate(self, docid: str) -> AsyncIterator:
        """
        A document mutation loop::

            async for doc in couch.mutate_doc("spam"):
                doc.foo = "bar"

        Will replay the mutation until it goes through.
        """
        doc = await self.get(docid)
        while True:
            yield doc
            try:
                await self.attempt_put(doc)
            except Conflict:
                doc = await self.get(docid)
            else:
                break

    async def iter_all_docs(
        self, *, include_docs: bool = False
    ) -> AsyncIterator[structs.AllDocs_DocRef]:
        """
        List all documents.

        This excludes design documents, see :meth:`.iter_design_docs`.

        Args:
            include_docs: Pre-load documents

        See :http:get:`/{db}/_all_docs`
        """
        # TODO: Pagination
        resp = await self._session._request(
            "GET",
            self._name,
            "_all_docs",
            params={
                "include_docs": include_docs,
            },
            headers={
                "Accept": "application/json",
            },
        )
        blob = resp.json()
        for ref in blob["rows"]:
            if ref["id"].startswith("_design/"):
                continue
            if "doc" in ref:
                doc = self._blob2doc(ref["doc"], self._name, ref["id"])
            else:
                doc = None
            yield structs.AllDocs_DocRef(
                _db=self, docid=ref["id"], rev=ref["value"]["rev"], _doc=doc
            )

    async def iter_indexes(self) -> AsyncIterator[structs.Index]:
        resp = await self._session._request("GET", self._name, "_index")
        payload = resp.json()
        for idx in payload["indexes"]:
            if idx["ddoc"] is None and idx["name"] == "_all_docs":
                continue
            yield structs.Index.from_dict(idx)

    async def add_index(
        self,
        name: str | None = None,
        *,
        fields: list[str],
        ddoc: str | None = None,
        type: str = "json",
    ):
        await self._session._request(
            "POST",
            self._name,
            "_index",
            json={
                "index": {"fields": list(fields)},
                "name": name,
                "ddoc": ddoc,
                "type": type,
            },
        )


class NoServerFound(Exception):
    """
    None of the configured servers seem to be working.
    """


class SessionPool:
    """
    Responsible for giving out Couch connections.

    You probably want to override like::

        class MyPool(SessionPool):
            session_class = MySession
    """

    _client: httpx.AsyncClient

    #: Class to use for sessions
    session_class: type[CouchSession]

    def __init__(self):
        super().__init__()
        self._client = self.make_client()

    def make_client(self) -> httpx.AsyncClient:
        """
        Produce an httpx client.
        """
        return httpx.AsyncClient(http2=True, follow_redirects=True)

    async def iter_servers(self) -> AsyncIterator[str | httpx.URL]:
        """
        Produce the list of potential servers.

        Override this
        """
        raise NotImplementedError
        for _ in ():
            yield

    async def _check_server(self, url: httpx.URL):
        resp = await self._client.get(url.join("_up"))
        resp.raise_for_status()

    async def session(self) -> CouchSession:
        """
        Get a session
        """
        excs = []
        async for url in self.iter_servers():
            url = httpx.URL(url)
            try:
                await self._check_server(url)
            except Exception as e:
                excs.append(e)
            else:
                return self.session_class(self._client, url)
        else:
            raise NoServerFound() from ExceptionGroup(
                "There were errors checking servers", excs
            )
