import base64
import json
import logging
import pickle
import sys
from importlib import import_module
from urllib.parse import parse_qs

from channels import Group
from channels.auth import channel_session_user,channel_session_user_from_http
from channels.exceptions import RequestAborted,RequestTimeout
from channels.generic.websockets import WebsocketConsumer
from channels.handler import AsgiRequest
from django import http
from django.conf import settings
from django.core import signals
from django.core.handlers import base
from django.http import FileResponse,HttpResponse
from django.urls import set_script_prefix

session_headers=['server','headers','client','query_string']

logger=logging.getLogger(__name__)

class WebSocketRequest(AsgiRequest):
  def __init__(self,message):
    for key in session_headers:
      message.content[key]=pickle.loads(base64.b64decode(str.encode(message.channel_session[key])))
    text=message.content["text"]

    obj=json.loads(text)
    path=obj.get("path")
    message["path"]=path
    if "?" in path:
      path,query_string=path.split('?',1)
      message.content['query_string']=query_string
    super(WebSocketRequest,self).__init__(message)
    self.POST=obj.get("data") or {}
    self.method=obj.get("method") or "GET"

class WebSocketHandler(base.BaseHandler):
  request_class=WebSocketRequest
  chunk_size=512*1024

  def __init__(self,*args,**kwargs):
    super(WebSocketHandler,self).__init__(*args,**kwargs)
    self.load_middleware()

  def __call__(self,message):
    set_script_prefix(message.get('root_path','') or '')
    signals.request_started.send(sender=self.__class__,message=message)
    try:
      request=self.request_class(message)
    except UnicodeDecodeError:
      logger.warning('Bad Request (UnicodeDecodeError)',exc_info=sys.exc_info(),extra={'status_code':400,})
      response=http.HttpResponseBadRequest()
    except RequestTimeout:
      response=HttpResponse("408 Request Timeout (upload too slow)",status=408)
    except RequestAborted:
      return
    else:
      try:
        response=self.get_response(request)
        if isinstance(response,FileResponse): response.block_size=1024*512
      except WebSocketRequest.ResponseLater:
        return

    return response

class WebSocketToHttpRouter(WebsocketConsumer):
  handler_class=WebSocketHandler

  def __init__(self,message,**kwargs):
    self.handler=self.handler_class()
    super(WebSocketToHttpRouter,self).__init__(message,**kwargs)

  http_user=True
  strict_ordering=False
  SessionStore=import_module(settings.SESSION_ENGINE).SessionStore

  def connect(self,message,**kwargs):
    for key in session_headers:
      message.channel_session[key]=base64.b64encode(pickle.dumps(message.content[key])).decode()
    self.message.reply_channel.send({"accept":True})

  def receive(self,text=None,bytes=None,**kwargs):
    response=self.handler(self.message)

    content=response.content.decode('utf-8')
    if 'application/json'==response['Content-Type']:
      content=json.loads(content)

    self.send(json.dumps({
      "statusCode":response.status_code,
      "content":content
    }))

  def disconnect(self,message,**kwargs):
    pass

# Connected to websocket.connect

#@channel_session
@channel_session_user_from_http
def ws_connect(message,room_name):
  message.reply_channel.send({"accept":True})
  params=parse_qs(message.content["query_string"])
  if b"username" in params:
    # Set the username in the session
    message.channel_session["username"]=params[b"username"][0].decode("utf8")
    Group("chat-%s"%room_name).add(message.reply_channel)
  else:
    message.reply_channel.send({"close":True})

# Connected to websocket.receive
#@channel_session
@channel_session_user
def ws_message(message,room_name):
  print(message.user)
  print('ws_message',room_name,message)
  Group("chat-%s"%room_name).send({
    "text":json.dumps({
      "text":message["text"],
      "username":message.channel_session["username"],
    }),
  })

# Connected to websocket.disconnect
#@channel_session
@channel_session_user
def ws_disconnect(message,room_name):
  print(message.user.username)
  Group("chat-%s"%room_name).discard(message.reply_channel)
