import asyncio as aio
import base64
import inspect
import json
import traceback
from functools import wraps
from itertools import count
from pathlib import Path
from uuid import uuid4 as uuid

from hrepr import Tag
from starlette.exceptions import HTTPException
from starlette.responses import (
    FileResponse,
    HTMLResponse,
    JSONResponse,
    RedirectResponse,
    Response,
)
from starlette.routing import Route, WebSocketRoute
from starlette.websockets import WebSocketDisconnect

from .page import Page
from .repr import Representer
from .utils import keyword_decorator

here = Path(__file__).parent

_count = count()

construct = {}


def register_constructor(key):
    def deco(fn):
        construct[key] = fn
        return fn

    return deco


@register_constructor("HTMLElement")
def _(page, selector):
    return page[selector]


class Queue2(aio.Queue):
    def putleft(self, entry):
        self._queue.appendleft(entry)
        self._unfinished_tasks += 1
        self._finished.clear()
        self._wakeup_next(self._getters)


class Cub:
    def __init__(self, mother, process, query_params={}, session={}):
        self.mother = mother
        self.fn = mother.fn
        self.process = process
        self.query_params = query_params
        self.session = session
        self.route = self.mother.path_for("main", process=self.process).rstrip("/")
        self.methods = {}
        self.representer = Representer(self.route)
        self.iq = aio.Queue()
        self.oq = Queue2()
        self.history = []
        self.reset = False
        self.ws = None
        self.page = Page(
            self.iq,
            self.oq,
            representer=self.representer,
            query_params=query_params,
            session=session,
        )
        self.coro = aio.create_task(self.run())
        self._sd_coro = None

    def schedule_selfdestruct(self):
        async def sd():
            await aio.sleep(self.mother.process_timeout)
            del self.mother.cubs[self.process]
            self.coro.cancel()
            print(f"Destroyed process: {self.process}")

        if self.mother.process_timeout is not None:
            self._sd_coro = aio.create_task(sd())

    def unschedule_selfdestruct(self):
        if self._sd_coro:
            self._sd_coro.cancel()
            self._sd_coro = None

    async def run(self):
        try:
            await self.fn(self.page)
            await self.page.sync()
        except Exception as exc:
            traceback.print_exc()

    def _object_hook(self, dct):
        if "%" in dct:
            args = dict(dct)
            args.pop("%")
            return construct[dct["%"]](self.page, **args)
        else:
            return dct

    async def json(self, request):
        body = await request.body()
        return json.loads(body, object_hook=self._object_hook)

    async def route_main(self, request):
        self.unschedule_selfdestruct()
        with open(here / "base-template.html") as tpf:
            self.reset = True
            return HTMLResponse(tpf.read().replace("{{{route}}}", self.route))

    async def route_socket(self, ws):
        self.unschedule_selfdestruct()

        async def recv():
            while True:
                try:
                    data = await ws.receive_json()
                    self.iq.put_nowait(data)
                except WebSocketDisconnect:
                    break
            print("recv stopped")

        async def send():
            while True:
                txt, in_history = await self.oq.get()
                try:
                    await ws.send_text(txt)
                    if in_history:
                        self.history.append(txt)
                except RuntimeError:
                    # Put the unsent element back into the queue
                    self.oq.putleft((txt, in_history))
                    break
            print("send stopped")

        if self.ws:
            try:
                await self.ws.close()
            except RuntimeError:
                pass

        await ws.accept()
        self.ws = ws

        if self.reset:
            for entry in self.history:
                await ws.send_text(entry)
            self.reset = False

        await aio.wait(
            [aio.create_task(recv()), aio.create_task(send())],
            return_when=aio.FIRST_COMPLETED,
        )
        self.schedule_selfdestruct()

    async def route_method(self, request):
        method_id = request.path_params["method"]
        method = self.representer.callback_registry.resolve(method_id)
        try:
            args = await self.json(request)
        except json.JSONDecodeError:
            args = [await request.body()]
        result = method(*args, **request.query_params)
        if inspect.iscoroutine(result):
            result = await result
        if isinstance(result, Tag):
            return HTMLResponse(self.representer(result))
        else:
            return JSONResponse(result)

    async def route_post(self, request):
        data = await self.json(request)
        if "value" in data:
            self.representer.future_registry.resolve(
                fid=data["reqid"],
                value=data["value"],
            )
        else:
            self.representer.future_registry.reject(
                fid=data["reqid"],
                error=data["error"],
            )
        return JSONResponse({"status": "ok"})

    async def route_queue(self, request):
        data = await self.json(request)
        self.representer.queue_registry.put(
            qid=data["reqid"],
            value=data["value"],
        )
        return JSONResponse({"status": "ok"})


def get_process_from_request(request):
    process_base = request.path_params.get("process", None)
    if process_base is None:
        process_base = base64.urlsafe_b64encode(uuid().bytes).decode("utf8").strip("=")
    return process_base


def forward_cub(fn):
    @wraps(fn)
    async def fwd(self, request):
        process = get_process_from_request(request)
        cub = self._get(process, query_params=request.query_params)
        if cub is None:
            print(f"Trying to access missing process: {process}")
            return JSONResponse({"missing": process}, status_code=404)
        else:
            return await fn(self, request, cub)

    return fwd


class MotherBear:
    def __init__(self, fn, process_timeout=60, hide_processes=True):
        self.fn = fn
        self.router = None
        self.process_timeout = process_timeout
        self.hide_processes = hide_processes
        self.cubs = {}
        self.appid = next(_count)

    def _get(self, proc, query_params={}, session={}, ensure=False):
        if proc not in self.cubs:
            if ensure:
                print(f"Creating process: {proc}")
                self.cubs[proc] = Cub(
                    self, proc, query_params=query_params, session=session
                )
            else:
                return None
        return self.cubs[proc]

    def path_for(self, name, **kwargs):
        return self.router.url_path_for(self._mangle(name), **kwargs)

    def _ensure_router(self, request):
        router = request.scope["router"]
        if self.router is None:
            self.router = router
        else:
            assert self.router is router

    async def route_dispatch(self, request):
        self._ensure_router(request)
        process = get_process_from_request(request)
        main_path = self.path_for("main", process=process)
        if self.hide_processes:
            try:
                session = request.session
            except AssertionError:
                session = {}
            return await self._get(
                process, query_params=request.query_params, session=session, ensure=True
            ).route_main(request)
        else:
            url = main_path
            if request.query_params:
                url = f"{url}?{request.query_params}"
            return RedirectResponse(url=url)

    async def route_main(self, request):
        self._ensure_router(request)
        process = get_process_from_request(request)
        return await self._get(
            process, query_params=request.query_params, ensure=True
        ).route_main(request)

    @forward_cub
    async def route_socket(self, ws, cub):
        self._ensure_router(ws)
        return await cub.route_socket(ws)

    @forward_cub
    async def route_method(self, request, cub):
        self._ensure_router(request)
        return await cub.route_method(request)

    @forward_cub
    async def route_file(self, request, cub):
        self._ensure_router(request)
        pth = cub.representer.file_registry.get_file_from_url(
            request.path_params["path"]
        )
        if pth is None:
            raise HTTPException(
                status_code=404, detail="File not found or not available."
            )
        return FileResponse(pth, headers={"Cache-Control": "no-cache"})

    @forward_cub
    async def route_vfile(self, request, cub):
        self._ensure_router(request)
        vf = cub.representer.vfile_registry.get(request.path_params["path"])
        return Response(content=vf.content, media_type=vf.type)

    @forward_cub
    async def route_post(self, request, cub):
        self._ensure_router(request)
        return await cub.route_post(request)

    @forward_cub
    async def route_queue(self, request, cub):
        self._ensure_router(request)
        return await cub.route_queue(request)

    async def route_static(self, request):
        self._ensure_router(request)
        pth = here / request.path_params["path"]
        return FileResponse(pth, headers={"Cache-Control": "no-cache"})

    def _mangle(self, name):
        return f"app{self.appid}_{name}"

    def _make_route(self, name, path, cls=Route, **kwargs):
        return cls(
            path,
            getattr(self, f"route_{name}"),
            name=self._mangle(name),
            **kwargs,
        )

    def routes(self):
        return [
            self._make_route("dispatch", "/"),
            self._make_route("static", "/{process:str}/static/{path:path}"),
            self._make_route("main", "/{process:str}/"),
            self._make_route(
                "method",
                "/{process:str}/method/{method:int}",
                methods=["GET", "POST"],
            ),
            self._make_route("file", "/{process:str}/file/{path:path}"),
            self._make_route("vfile", "/{process:str}/vfile/{path:path}"),
            self._make_route("post", "/{process:str}/post", methods=["POST"]),
            self._make_route("queue", "/{process:str}/queue", methods=["POST"]),
            self._make_route("socket", "/{process:str}/socket", cls=WebSocketRoute),
        ]


@keyword_decorator
def bear(fn, **kwargs):
    return MotherBear(fn, **kwargs)
