import os
import hashlib
import logging
import webbrowser
from textwrap import dedent
from datetime import datetime

from flask import Flask, request, render_template_string, send_file
from geventwebsocket.handler import WebSocketHandler
from werkzeug.serving import run_with_reloader
from gevent.pywsgi import WSGIServer
from gevent import queue
import gevent

from akindi.utils import extend, json
from akindi.imgproc.ccimgproc_compat import result_to_cv

log = logging.getLogger(__name__)

index_template = dedent("""\
    <!DOCTYPE HTML>
    <html>
    <head>
        <title>Proctor Server UI</title>
        <script>
            window.PT_WS_URL = "ws://" + window.location.host + "/ws/client";
        </script>
        <script type="text/javascript" src="/static/imgproc.js"></script>
    </head>
    <body ng-app="akindi.imgproc.proctor">
        <imgproc-proctor>
            Loading&hellip;
        </imgproc-proctor>
    </body>
    </html>
""")

app = Flask(__name__)
env = None
server_bind = None
active_server = None
clients = []

def file_sha1(path):
    sha1 = hashlib.sha1()
    with open(path) as f:
        for chunk in iter(lambda: f.read(4096), ""):
            sha1.update(chunk)
    return sha1.hexdigest()

class WSBase(object):
    def _handle_msg(self, msg_str):
        try:
            msg = json.loads(msg_str)
        except ValueError as e:
            log.warning("Error parsing message %r: %r (ignoring)", msg_str, e)
            return
        type_name = msg.get("type", "<no-type>").replace("-", "_")
        handler = getattr(self, "HANDLE_%s" %(type_name, ), None)
        if not handler:
            log.warning("No handler for message with type %r (ignoring)", type_name)
            return

        msg_data = msg.get("data")
        if msg_data is None:
            log.warning("No data in message %r (ignoring)", msg)
            return
        return handler(msg_data)


class WSClientHandler(WSBase):
    def __init__(self, env, active_server, ws):
        self.db = env.db
        self.env = env
        self.ws = ws
        self.active_server = active_server
        self._send_queue = queue.Queue(maxsize=10)
        self._sender = gevent.spawn(self._send_thread)
        self._init()

    def _init(self):
        if self.active_server:
            self.init_from_server(self.active_server)
        else:
            self.init_from_db()
        clients.append(self)

    def shutdown(self):
        if self in clients:
            clients.remove(self)
        self._sender.kill()
        self.ws.close()

    def init_from_server(self, server):
        for msg in server.messages:
            self.send(msg)

    def init_from_db(self):
        db = self.db
        run = (
            db.session.query(db.ImgprocTestRun)
                .order_by(db.ImgprocTestRun.id.desc())
                .first()
        )
        if not run:
            return
        self.send(run.get_message_json())
        seen_pdfs = set()
        for page in run.pages:
            if page.pdf_sha1 not in seen_pdfs:
                seen_pdfs.add(page.pdf_sha1)
                pdf = (
                    db.session.query(db.ImgprocPDF)
                        .filter_by(sha1=page.pdf_sha1)
                        .first()
                )
                if not pdf:
                    log.warning("No PDF found for page %s!", page)
                else:
                    self.send(pdf.get_message_json())
            self.send(page.get_message_json())

    def _send_thread(self):
        while self._send_queue.peek() or True:
            msgs = []
            while len(msgs) < 10 and not self._send_queue.empty():
                msgs.append(self._send_queue.get_nowait())
            msg_str = json.dumps(msgs)
            self.ws.send(msg_str)

    def send(self, msg):
        self._send_queue.put(msg)
    
    def send_nowait(self, msg):
        self._send_queue.put_nowait(msg)

    def run(self):
        while True:
            msg_str = self.ws.receive()
            if msg_str is None:
                break
            self._handle_msg(msg_str)

    def HANDLE_send_oracle(self, msg):
        db = self.env.db
        oracles = (
            db.session.query(db.PDFPageOracle)
                .filter(db.PDFPageOracle.imgproc_pages.any(test_run_id=4))
                .filter_by(name=msg["name"])
                .order_by(db.PDFPageOracle.id.asc())
        )
        for oracle in oracles:
            self.send(oracle.get_message_json())


@app.route("/ws/client")
def ws_client_view():
    ws = request.environ.get("wsgi.websocket")
    if ws is None:
        return "Error: must be a websocket request"
    print "Client connected"
    try:
        client = WSClientHandler(env, active_server, ws)
        client.run()
    finally:
        print "Client lost"
        client.shutdown()
    return "DONE"


class WSServerHandler(WSBase):
    def __init__(self, env, ws):
        self.db = env.db
        self.env = env
        self.ws = ws
        self.messages = []

    def run(self):
        self.test_run = None
        while True:
            msg_str = self.ws.receive()
            if msg_str is None:
                self._cleanup()
                return
            res = self._handle_msg(msg_str)
            if res is None:
                continue
            self.db.session.flush()
            message = res.get_message_json()
            self.messages.append(message)
            for client in clients:
                try:
                    client.send_nowait(message)
                except queue.Full:
                    print "Discarding client with full queue"
                    client.shutdown()

    def _cleanup(self):
        if self.test_run:
            self.test_run.finished_on = datetime.now()
        self.db.session.flush()

    def HANDLE_run_start(self, msg):
        if not clients:
            webbrowser.open("http://%s:%s/" %server_bind)
        self.test_run = self.db.ImgprocTestRun.create(
            argv=json.dumps(msg["argv"]),
            started_on=datetime.now(),
        )
        return self.test_run

    def HANDLE_pdf_start(self, msg):
        db = self.db
        if "path" not in msg:
            raise Exception("Got pdf-start without path: %r" %(msg, ))
        if "name" not in msg:
            msg["name"] = os.path.basename(msg["path"])
        if "sha1" not in msg:
            msg["sha1"] = file_sha1(msg["path"])
        kwds = dict((k, msg[k]) for k in ["sha1", "name", "path"])
        pdf = db.session.query(db.ImgprocPDF).filter_by(**kwds).first()
        if not pdf:
            pdf = self.db.ImgprocPDF.create(**kwds)
        self.last_pdf_sha1 = kwds["sha1"]
        return pdf

    def HANDLE_page_add(self, msg):
        return self.db.ImgprocPage.create(
            test_run=self.test_run,
            pdf_sha1=msg["pdf_sha1"],
            page_num=msg["page_num"],
            cvpage=json.dumps(msg["cvpage"]),
        )

    def HANDLE_page_add_ccimgproc(self, msg):
        cvpage = extend(msg["result"], {
            "cv": result_to_cv(msg["result"]),
            "images": msg["images"],
            "page_num": msg["page_num"],
            "debug": msg["result"].pop("debug_info", {}),
        })
        cvpage.pop("bubbles", None)
        return self.db.ImgprocPage.create(
            test_run=self.test_run,
            pdf_sha1=self.last_pdf_sha1,
            page_num=cvpage["page_num"],
            cvpage=json.dumps(cvpage),
        )

@app.route("/ws/server")
def ws_server_view():
    global active_server
    ws = request.environ.get("wsgi.websocket")
    if ws is None:
        return "Error: must be a websocket request"
    if active_server is not None:
        log.error("attempt to have multiple active servers")
        return "Error: Only one server can be active at a time"
    try:
        active_server = WSServerHandler(env, ws)
        active_server.run()
    finally:
        active_server = None
    return "ok"


@app.route("/")
def index_view():
    return render_template_string(index_template)


@app.route("/api/unsafe-get-page")
def api_unsafe_get_page():
    path = request.args.get("page")
    if not path:
        return "ERROR: no path"
    if not path.endswith(".png"):
        return "ERROR: invalid path"
    return send_file(path, cache_timeout=20)


def run(env_, config):
    global env, server_bind
    default_config = {
        "ws_host": "127.0.0.1",
        "ws_port": 4410,
        "static_folder": None, # will be akindi/media/ui-angular/dist/
        "reloader": True,
    }

    env = env_

    config = extend(default_config, config)
    app.static_folder = config["static_folder"]
    server_bind = (config["ws_host"], config["ws_port"])
    app.debug = True

    def run_server():
        server = WSGIServer(
            server_bind, app, log=log,
            handler_class=WebSocketHandler,
        )
        server.logger = log
        print "Starting proctor server on http://%s:%s/" %server_bind
        server.serve_forever()
    if config["reloader"]:
        run_server = run_with_reloader(run_server)
    run_server and run_server()
