import dataclasses
import functools
import threading
import typing

from fastapi.responses import UJSONResponse
from starlette.responses import HTMLResponse
from termcolor import colored
import mmcore
from mmcore.base import AGroup
from mmcore.base.params import pgraph
from mmcore.base.userdata.controls import find_points_in_controls
from mmcore.common.viewer import control_points_observer
from mmcore.gql.lang.parse import parse_simple_query

TYPE_CHECKING = False
import ujson as json
import strawberry
from mmcore.collections.multi_description import ElementSequence
from strawberry.scalars import JSON
import uvicorn, fastapi
from mmcore.base.registry import *
from rpyc import ThreadPoolServer, SlaveService
from mmcore.geom.point import BUFFERS
import IPython
from mmcore.base.tags import __databases__

DBs = __databases__

__node_resolvers__ = dict()
def search_all_indices(lst, value):
    for i, v in enumerate(lst):
        if v == value:
            yield i


def generate_uvicorn_app_name(fpath, appname="app"):
    if "\\" in fpath:
        r = list(fpath.split("\\"))
    else:
        r = list(fpath.split("/"))
    if r[-1].startswith("__"):
        r = r[:-1]

    return ".".join(r[list(search_all_indices(r, "mmcore"))[-1]:]) + ":" + appname


class AllDesc:
    def __init__(self, default=None):
        self.default = default

    def __set_name__(self, owner, name):
        self.name = name

    def __get__(self, instance, owner):
        if not instance:
            return self.default
        return json.loads(instance.ToJSON())


_server_binds = []
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware

from fastapi import WebSocket
import mmcore.base.models.gql as gql_models
from graphql.language import parse
import graphql

from graphene import ObjectType, String, Schema

rpyc_input = dict()

rpyc_namespace = dict(post=lambda k, v: rpyc_input.__setitem__(json.dumps(k), json.loads(v)))

TYPE_CHECKING = False
# language=GraphQl
sch2 = """
schema {
    query: query_root
    mutation: mut_root
}

input ObjectInput {
    name:String


}
scalar JSON
type UserData {
    properties: JSON
}



type Object3D {
    all: JSON
    castShadow: Boolean
    children : [Object3D]
    layers: Int!
    matrix: [Float!]!
    name: String
    receiveShadow: Boolean
    type: String!
    up: [Float!]
    userData: UserData
    uuid: String!
    geometry: String}
type query_root{
    root:Root
}
type Root {
    geometries: [JSON]
    materials: [JSON]
    metadata: JSON
    object: Object3D
}

type mut_root{
    root: Root
}


"""

sch22 = parse(sch2)
sch3 = graphql.build_ast_schema(sch22)
ee = ElementSequence(list(objdict.values()))

T = typing.TypeVar("T")

from starlette.exceptions import HTTPException

from mmcore.gql.server.fastapi import MmGraphQlAPI

debug_properties = {}


@dataclasses.dataclass
class ConsoleStartMessage:
    text: str = f"\n    \u002B\u00D7  mmcore {mmcore.__version__()}\n"
    color: str = 'light_magenta'
    on_color: typing.Optional[str] = None
    attrs: typing.Optional[tuple[str]] = ('bold',)


@dataclasses.dataclass
class PropsUpdate:
    uuids: list[str]
    props: dict[str, typing.Any]


class IDict(dict):
    def __init__(self, dct=idict, table=adict):
        super().__init__(dct)
        self.table = table
        self.dct = dct

    def __getitem__(self, k):
        return self.trav(k)

    def trav(self, i):
        print(self.table[i])
        if 'AGroup' == self.table[i].__class__.__name__:
            return {'kind': "Group", "items": [self.trav(j) for j in list(self.dct[i]["__children__"])]}
        else:
            obj = {'kind': "Object", "name": i, "value": self.table[i]}
            if self.dct[i] == {}:
                obj |= {"isleaf": True}
                return obj

            elif isinstance(self.dct[i], (list, tuple, set)):
                if isinstance(self.dct[i], set):
                    dct = list(self.dct[i])
                else:
                    dct = self.dct[i]
                if len(dct) == 0:
                    obj |= {'kind': "Group", "isleaf": True}
                    return obj
                else:
                    obj |= {'kind': "Group", "items": [self.trav(j) for j in dct], "isleaf": False}
                    return obj

            else:
                _dct = {}
                for k, v in self.dct[i].items():
                    _dct[k] = self.trav(v)
                obj |= _dct
                obj |= {"isleaf": False}
                return VDict(obj)


class Query(ObjectType):
    name = String()


schema = Schema(Query)


def _create_fastapi(name, mounts=(), mws=(
        (
                CORSMiddleware,
                dict(allow_origins=["*"],
                     allow_methods=["GET", "POST", "PUT", "HEAD", "OPTIONS", "DELETE"],
                     allow_headers=["*"],
                     allow_credentials=["*"])),
        (GZipMiddleware, dict(minimum_size=500))
), **kwargs):
    _ = fastapi.FastAPI(name=name, **kwargs)
    if len(mounts) > 0:
        for path, v, mntname in iter(mounts):
            _.mount(path, v, mntname)
    if len(mws) > 0:
        for obj, props in iter(mws):
            _.add_middleware(obj, **props)
    return _


class VDict(dict):
    def __init__(self, *args, relay=idict, table=adict, **kwargs):
        dict.__init__(self, *args, **kwargs)
        self.dct = relay
        self.table = table

    def __getitem__(self, item):
        if item == "value":
            return dict.__getitem__(self, item).value()

        else:
            return dict.__getitem__(self, item)


@functools.lru_cache(maxsize=None)
def find_parent_groups(obj):
    vv = []
    for k, v in adict.items():
        if isinstance(v, AGroup):
            if obj in v:
                if v.uuid not in vv:
                    vv.append(v.uuid)

            elif obj in v.children:
                vv.extend(find_parent_groups(v))
    return [adict[i] for i in vv]

class SharedStateServer():

    port: int = 7711
    rpyc_port: int = 7799
    host: str = "0.0.0.0"
    appname: str = "serve_app"
    prodapp = None
    _is_production = False
    console_start_message: ConsoleStartMessage = ConsoleStartMessage()
    @property
    def is_production(self):
        return self._is_production

    @is_production.setter
    def is_production(self, v):
        self._is_production = v

    def __new__(cls, *args, header="mmcore api", **kwargs):
        global serve_app, DBs

        import threading as th
        inst = object.__new__(cls)
        inst.__dict__ |= kwargs
        inst.header = header

        gqlv2app = MmGraphQlAPI(gql_endpoint="/graphql")

        @gqlv2app.post(gqlv2app.gql_endpoint)
        def graphql_query_resolver(data: dict):
            ##print(data)

            qt2 = parse_simple_query(data['query'])

            return qt2.resolve(pgraph)

        _server_binds.append(inst)

        def shutdown():
            print("shutdown event...")
            for db in DBs.values():
                db.save()

        serve_app = fastapi.FastAPI(name='serve_app', title=inst.header, on_shutdown=[shutdown])

        inst.app = serve_app
        serve_app.mount("/v2", gqlv2app, "gqlv2app")
        inst.app.add_middleware(CORSMiddleware, allow_origins=["*"],
                                allow_methods=["GET", "POST", "PUT", "HEAD", "OPTIONS", "DELETE"],
                                allow_headers=["*"],
                                allow_credentials=["*"])
        inst.app.add_middleware(GZipMiddleware, minimum_size=500)
        inst.resolvers = dict()
        inst.jobs = dict()
        @strawberry.type
        class Root(typing.Generic[T]):
            __annotations__ = {
                "all": typing.Optional[JSON],
                "shapes": typing.Optional[JSON],
                "metadata": gql_models.Metadata,
                "materials": list[gql_models.AnyMaterial],
                "object": T,
                "geometries": list[typing.Union[gql_models.BufferGeometry, None]]
            }

        def pull():
            br = objdict["_"]
            for o in objdict.values():
                if not o.name == "base_root":
                    br.add(o)
            return br

        @inst.app.get("/buffers/{uid}")
        def get_buffer(uid: str):
            print(f"GET {uid} buffer {BUFFERS[uid]._buffer}")

            return {"data": BUFFERS[uid]._buffer}

        @inst.app.post("/buffers/{uid}")
        def upd_all_buffer(uid: str, data: dict):

            BUFFERS[uid].update_all(data["data"])

            return {"data": BUFFERS[uid]._buffer}

        @inst.app.post("/buffers/{uid}/{index}")
        def upd_item_buffer(uid: str, index: int, data: dict):
            if not inst.is_production:
                print(f"[POST] update index {index} for '{uid}' buffer: {BUFFERS[uid][index]}->{data['data']}>",
                      )

            BUFFERS[uid].update_item(index, data["data"])

            return {"data": BUFFERS[uid]._buffer}

        @inst.app.get("/test/app/{uid}")
        def appp(uid: str):
            # language=JavaScript
            application = """
            const socket = new WebSocket('ws://${host}:${port}/ws');
            socket.addEventListener('open', function (event) {socket.send('{ "uuid":"${uid}" }');});
            socket.addEventListener('message', function (event) {console.log(event.data);});
            socket.onmessage
            """.replace('${host}', str(inst.host)).replace('${port}', str(inst.port)).replace('${uid}', uid)
            # language=Html
            return HTMLResponse("""
        <!DOCTYPE html>
        <body>
            <code>
                <script>
                    {script}
                </script>
            </code>
        </body>
            """.format(
                script=application
            )
            )

        @inst.app.get("/h")
        async def home2():
            return strawberry.asdict(pull()())

        @inst.app.get("/fetch/{uid}")
        async def get_item(uid: str):
            try:

                return adict[uid].root()
            except KeyError as errr:
                return HTTPException(401, detail=f"KeyError. Trace: {errr}")

        @inst.app.post("/props-update/{uid}")
        async def props_update(uid: str, data: PropsUpdate):
            try:
                target = adict[uid]

            except KeyError as errr:
                return HTTPException(401, detail=f"KeyError. Trace: {errr}")
            print(uid, data, target)

            target.props_update(data.uuids, data.props)
            return {"uuid": target.uuid}

        @inst.app.post("/fetch/{uid}")
        def post_item(uid: str, data: dict):
            if len(data) > 0:
                if pgraph.item_table.get(uid) is not None:
                    pgraph.item_table.get(uid)(**kwargs)

                else:
                    adict[uid](**data)

        @inst.app.get("/gui/{uid}")
        def gui_get_item(uid: str):

            return adict[uid].gui_get()

        @inst.app.post("/gui/{uid}")
        def gui_post_item(uid: str, data: dict):

            return adict[uid].gui_post(data)

        @inst.app.post("/controls/{uid}")
        def controls_post_item(uid: str, data: dict):

            control_points_observer.notify(uid, find_points_in_controls(data))
            return adict[uid].controls.todict()

        @inst.app.post("/controls/{uid}")
        def controls_get_item(uid: str):

            return adict[uid].controls.data.todict()

        @inst.app.get("/keys", response_model_exclude_none=True)
        async def keys():
            return list(adict.keys())

        @inst.app.get("/", response_model_exclude_none=True)
        async def home():
            return {
                'app': 'mmcore', 'version': str(mmcore.__version__())
                }

        @inst.app.post("/graphql")
        async def gql(data: dict):
            if "target" in debug_properties:

                return UJSONResponse(adict[debug_properties["target"]].root())
            else:
                qq = parse_simple_query(data["query"])

                return qq.resolve(adict)

        @inst.app.options("/graphql")
        async def gql():

            return "OK"

        @inst.app.get("/graphql")
        async def gql():

            return "OK"



        @inst.app.post("/", response_model_exclude_none=True)
        async def mutate(data: dict = None):
            if "target" in debug_properties:

                return adict[debug_properties["target"]].root()
            else:
                return {}

        @inst.app.post("/jobs/{name}", response_model_exclude_none=True)
        async def jobs_post(name: str, data: dict = None):
            inst.jobs[name](data)
            return {name: "OK"}




        @inst.app.websocket("/ws")
        async def websocket_endpoint(websocket: WebSocket):
            """

            @param websocket:
            @return:
            {
                # POST
                uuid: "some-item-uuid",
                body: {
                    # properties
                    name: "A"
                    ...
                }
                ---
                # GET
                uuid: "some-item-uuid",

            }
            """
            await websocket.accept()
            while True:
                data = await websocket.receive_json()
                # print(f"WS: {data}")
                obj = adict[data["uuid"]]
                if "body" in data.keys():
                    if not ((data["body"] is None) or (data["body"] == {})):
                        obj(**data["body"])

                await websocket.send_json(data=obj.root())

        def run_rpyc():
            service = SlaveService()
            service.namespace = rpyc_namespace
            _serv = ThreadPoolServer(service, port=cls.rpyc_port)
            _serv.start()

        inst.thread = th.Thread(target=inst.run)
        inst.rpyc_thread = th.Thread(target=run_rpyc)
        inst.runtime_env = dict(inputs=dict(), out=dict())
        inst.resolvers = dict()
        inst.params_nodes = dict()
        inst.loglevel = 'error'
        @inst.app.get("/ipykernel/connection_file")
        async def ipykernel_connection_file():
            data = inst.get_ipy_connection_file()
            if data is None:
                return {"data": None, "reason": "Kernel is not initialized"}
            return {"data": data, "reason": {}}

        @inst.app.post("/resolver/{name}/{uid}")
        async def external_post(name: str, uid: str, data: dict):
            return UJSONResponse(inst.resolvers[name](uid, data))

        @inst.app.get("/resolver/{name}/{uid}")
        async def external_get(name: str, uid: str):
            return UJSONResponse(inst.resolvers[name](uid))

        @inst.app.post("/params/node/{uid}")
        async def params_post(uid: str, data: dict):
            pgraph.item_table.get(uid)(**data)
            return pgraph.item_table.get(uid).todict(no_attrs=True)

        @inst.app.get("/params/node/{uid}")
        async def params_get(uid: str):
            return pgraph.item_table.get(uid).todict(no_attrs=True)

        @inst.app.get("/params/nodes")
        async def params_nodes_names():
            return list(pgraph.item_table.keys())

        @inst.app.get("/resolvers")
        def resolvers():
            return list(inst.resolvers.keys())

        @inst.app.get("/debug/props")
        async def debug_props():
            return debug_properties

        @inst.app.get("/flow")
        async def get_flow():
            return pgraph.toflow()

        if STARTUP:
            try:
                inst.thread.start()
            except OSError as err:
                print("Shared State server is already to startup. Shutdown...")
                raise err

        return inst

    def run(self, **kwargs):
        self.uvicorn_appname = generate_uvicorn_app_name(__file__, appname=self.appname)

        self.use_ssl = False
        for k in kwargs.keys():
            if k.startswith("ssl"):
                self.use_ssl = True
                break

        self.print()
        uvicorn.run(self.uvicorn_appname, port=self.port, host=self.host, log_level=self.loglevel, **kwargs)

    def print(self):
        _prefix = 'https' if    self.use_ssl else 'http'
        _print_host = self.host if self.host != '0.0.0.0' else 'localhost'
        print(colored(f"\n    \u002B\u00D7  mmcore {mmcore.__version__()}\n", 'light_magenta', None, ('bold',)))
        print(f'        server: uvicorn\n        app: {self.uvicorn_appname}\n        local: {_prefix}://{_print_host}'
              f':{self.port}/\n        openapi ui: {_prefix}://{_print_host}'
              f':{self.port}/docs\n'
              )


    def production_run(self, name, api_prefix, mounts=(), middleware=None, app_kwargs=None, **kwargs):
        if app_kwargs is None:
            app_kwargs = dict()

        if middleware is None:
            middleware = []
        for obj in self.app.user_middleware:
            if not (obj in [mn[0] for mn in middleware]):
                middleware.append(
                    (obj.cls, obj.options))

        self.prodapp = _create_fastapi(name,
                                       mounts=[
                                           (api_prefix, self.app, self.app.extra['name']),
                                           *mounts],
                                       mws=middleware, **app_kwargs
                                       )
        self.uvicorn_appname = generate_uvicorn_app_name(__file__, appname=name)
        print(f'running uvicorn {  self.uvicorn_appname}')
        self.is_production = True
        uvicorn.run(  self.uvicorn_appname, port=self.port, host=self.host, **kwargs)

    def stop(self):
        self.thread.join(6)

    def gltf_system(self):
        ...

    def create_child(self, path, name=None, title=None, **kwargs):
        if name is None:
            name = path.split("/")[-1]
        new_app = fastapi.FastAPI(name=name, title=self.header if title is None else title, **kwargs)

        new_app.add_middleware(CORSMiddleware, allow_origins=["*"],
                               allow_methods=["GET", "POST", "PUT", "HEAD", "OPTIONS", "DELETE"],
                               allow_headers=["*"],
                               allow_credentials=["*"])
        new_app.add_middleware(GZipMiddleware, minimum_size=500)

        self.app.mount(path, new_app, name)
        setattr(self, name, new_app)
        return new_app

    def start(self, port=None, host=None, log_level=None, **kwargs):

        for k, v in dict(port=port, host=host, log_level=log_level).items():
            if v is not None:
                setattr(self, k, v)
        self.thread = threading.Thread(target=self.run, args=(), kwargs=kwargs)
        self.thread.start()


    def stop_rpyc(self):
        self.rpyc_thread.join(6)

    def resolver(self, func):
        self.runtime_env["inputs"][func.__name__] = dict()
        self.runtime_env["out"][func.__name__] = dict()

        @functools.wraps(func)
        def wrapper(**kwargs):
            self.runtime_env["inputs"][func.__name__] = kwargs
            self.runtime_env["out"][func.__name__] = func(**self.runtime_env["inputs"][func.__name__])
            return self.runtime_env["out"][func.__name__]

        self.resolvers[str(func.__name__)] = wrapper
        return wrapper

    def add_resolver(self, name, func):
        @functools.wraps(func)
        def wrapper(**kwargs):
            self.runtime_env["inputs"][name] |= kwargs
            self.runtime_env["out"][name] |= func(**self.runtime_env["inputs"][name])
            return self.runtime_env["out"][name]

        self.resolvers[str(name)] = wrapper
        self.resolvers[str(name) + "__wrapped__"] = func
        return wrapper

    def add_params_node(self, name, fun):
        self.params_nodes[name] = fun

    def params_node(self, fun):
        if hasattr(fun, 'name'):
            self.params_nodes[fun.name] = fun
        else:
            self.params_nodes[fun.__name__] = fun
        return fun

    def bind(self, fun):
        def wrap(node):
            def call():
                fun(node)


    def start_rpyc(self):
        self.rpyc_thread.start()

    def run_thread(self):
        self.thread.run()

    def run_rpyc(self):
        self.rpyc_thread.run()

    def is_alive(self):
        return self.thread.is_alive()

    def is_alive_rpyc(self):
        return self.rpyc_thread.is_alive()

    def start_as_main(self, on_start=None, **kwargs):

        self.__dict__ |= kwargs
        if on_start:
            on_start()
        self.is_production = True

        self.run()

    def mount(self, path, other_app, name: str):
        self.app.mount(path, other_app, name)

    def event(self, fun):
        self.app.on_event()

    def start_ipython(self, argv=(), embed=True, prod=False, askernel=False, **kwargs):
        self.is_production = prod
        if askernel and embed:
            IPython.embed_kernel(header=self.header, **kwargs)
        elif embed:
            IPython.embed(header=self.header)

        else:
            IPython.start_ipython(argv=argv, **kwargs)

    def ipy(self):
        try:
            return eval("get_ipython()")
        except NameError as err:
            return None


serve = SharedStateServer()
