import os
import argparse
import socket
import tempfile

from avalon import io, api
from launcher import lib


def wait_for_maya_boot():
    # Create a TCP/IP socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    # Bind the socket to the port
    server_address = ("localhost", 20000)
    print("Starting up on {}".format(server_address))
    sock.bind(server_address)

    sock.settimeout(120)

    # Listen for incoming connections
    sock.listen(1)

    while True:
        # Wait for a connection
        print("Waiting for a connection to Maya.")
        connection, client_address = sock.accept()
        break


def send_to_maya(cmd, connection):
    print("Sending: {}".format(repr(cmd)))
    connection.send(cmd.encode())
    data = connection.recv(4096)
    print(data)


def read_maya_script_editor_output(maya_logging_file, previous_log):
    with open(maya_logging_file) as f:
        log = f.read().replace(previous_log, "")
        print(log)
        previous_log += log

    return previous_log


def main(silo, shots, task, application_name):
    # Maya logging.
    maya_logging_directory = tempfile.mkdtemp()
    maya_logging_file = os.path.join(
        maya_logging_directory, "maya_script_editor_output.txt"
    )
    os.environ["MAYA_CMD_FILE_OUTPUT"] = maya_logging_file
    previous_log = ""

    io.install()

    apps = lib.get_apps(io.find_one({"type": "project"}))
    application = None
    for App in apps:
        if App.name == application_name:
            application = App()

    session = api.Session
    session["AVALON_SILO"] = silo
    session["AVALON_TASK"] = task
    session["WORKFILES_STARTUP"] = ""
    os.environ["PYTHONPATH"] += os.pathsep + os.path.dirname(__file__)
    for shot in shots:
        session["AVALON_ASSET"] = shot
        session["AVALON_HIERARCHY"] = "/".join(
            io.find_one({"type": "asset", "name": shot})["data"]["parents"]
        )
        template = io.find_one({"type": "project"})["config"]["template"]
        session["AVALON_WORKDIR"] = template["work"].format(
            root=session["AVALON_PROJECTS"],
            project={"name": session["AVALON_PROJECT"]},
            hierarchy=session["AVALON_HIERARCHY"],
            asset=shot,
            task=task
        )

        application.process(session.copy())

        print("Waiting until maya is ready to go.")
        wait_for_maya_boot()
        previous_log = read_maya_script_editor_output(
            maya_logging_file, previous_log
        )

        # Establish connection.
        connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        connection.connect(("localhost", 7777))

        # Ensure no unsaved changes dialog appears.
        cmd = "from maya import cmds;cmds.file(modified=False)"
        send_to_maya(cmd, connection)
        previous_log = read_maya_script_editor_output(
            maya_logging_file, previous_log
        )

        # Open workfile.
        cmd = (
            "from avalon.tools import workfiles;"
            "from avalon import api;"
            "host = api.registered_host();"
            "window = workfiles.app.Window(host.work_root());"
            "window.on_open_pressed()"
        )
        send_to_maya(cmd, connection)
        previous_log = read_maya_script_editor_output(
            maya_logging_file, previous_log
        )

        # Ensure imageplane refreshes
        cmd = (
            "import pymel.core as pm;"
            "pm.lookThru(pm.PyNode(\"reviewMain\").members()[0])"
        )
        send_to_maya(cmd, connection)
        previous_log = read_maya_script_editor_output(
            maya_logging_file, previous_log
        )

        # Publish.
        cmd = "import pyblish.util;pyblish.util.publish()"
        send_to_maya(cmd, connection)
        previous_log = read_maya_script_editor_output(
            maya_logging_file, previous_log
        )

        # Close Maya
        cmd = "from maya import cmds;cmds.quit(force=True)"
        send_to_maya(cmd, connection)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Batch publish.")
    parser.add_argument(
        "--silo",
        dest="silo",
        action="store",
        required=True,
        help="Silo to find shots."
    )
    parser.add_argument(
        "--shots",
        dest="shots",
        nargs="+",
        required=True,
        help="Shots to publish."
    )
    parser.add_argument(
        "--task",
        dest="task",
        action="store",
        required=True,
        help="Task to publish."
    )
    parser.add_argument(
        "--application",
        dest="application",
        action="store",
        required=True,
        help="Task to publish."
    )

    args = vars(parser.parse_args())
    main(args["silo"], args["shots"], args["task"], args["application"])