#!/usr/bin/env python3
"""Simple tool to work with protobuf in pyatv."""

import argparse
import binascii
from collections import namedtuple
import difflib
import glob
from io import BytesIO
import os
import re
import stat
import subprocess
import sys
import zipfile

import cryptography
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from google.protobuf.text_format import MessageToString
import requests

# New messages re-using inner message of another type
REUSED_MESSAGES = {"DEVICE_INFO_MESSAGE": "DEVICE_INFO_UPDATE_MESSAGE"}

BASE_PATH = os.path.join("pyatv", "protocols", "mrp", "protobuf")
OUTPUT_TEMPLATE = """\"\"\"Simplified extension handling for protobuf messages.

THIS CODE IS AUTO-GENERATED - DO NOT EDIT!!!
\"\"\"

from .ProtocolMessage_pb2 import ProtocolMessage


{packages}


{messages}


_EXTENSION_LOOKUP = {{
    {extensions}
}}


{constants}


def _inner_message(self):
    extension = _EXTENSION_LOOKUP.get(self.type, None)
    if extension:
        return self.Extensions[extension]

    raise Exception('unknown type: ' + str(self.type))


ProtocolMessage.inner = _inner_message  # type: ignore
"""

MessageInfo = namedtuple("MessageInfo", ["module", "title", "accessor", "const"])


def _protobuf_url(version):
    base_url = (
        "https://github.com/protocolbuffers/protobuf/"
        + "releases/download/v{version}/protoc-{version}-{platform}.zip"
    )
    platforms = {
        "linux": "linux-x86_64",
        "darwin": "osx-x86_64",
        "win32": "win64",
    }

    platform = platforms.get(sys.platform)
    if not platform:
        print("Unsupported platform: " + sys.platform, file=sys.stderr)
        sys.exit(1)

    return base_url.format(version=version, platform=platform)


def _get_protobuf_version():
    with open("base_versions.txt", encoding="utf-8") as file:
        for line in file:
            match = re.match(r"protobuf==(\d+\.\d+\.\d+)[^0-9,]*", line)
            if match:
                return match.group(1)
    raise Exception("failed to determine protobuf version")


def _download_protoc(force=False):
    if os.path.exists(protoc_path()) and not force:
        print("Not downloading protoc (already exists)")
        return

    version = _get_protobuf_version()
    url = _protobuf_url(version)

    print("Downloading", url)

    resp = requests.get(url)
    with zipfile.ZipFile(BytesIO(resp.content)) as zip_file:
        for zip_info in zip_file.infolist():
            if zip_info.filename.startswith("bin/protoc"):
                print("Extracting", zip_info.filename)
                basename, extension = os.path.splitext(zip_info.filename)
                zip_info.filename = f"{basename}-{version}{extension}"
                zip_file.extract(zip_info)
                break

    if not os.path.exists(protoc_path()):
        print(protoc_path(), "was not downloaded correctly", file=sys.stderr)
        sys.exit(1)

    file_stat = os.stat(protoc_path())
    os.chmod(protoc_path(), file_stat.st_mode | stat.S_IEXEC)


def _verify_protoc_version():
    expected_version = _get_protobuf_version()
    try:
        ret = subprocess.run(
            [protoc_path(), "--version"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            check=False,
        )
        installed_version = ret.stdout.decode("utf-8").split(" ")[1].rstrip()
        if installed_version != expected_version:
            print(
                "Expected protobuf",
                expected_version,
                "but found",
                installed_version,
                file=sys.stderr,
            )
            sys.exit(1)
    except FileNotFoundError:
        print(
            "Protbuf compiler (protoc) not found. Re-run with --download",
            file=sys.stderr,
        )
        sys.exit(1)
    else:
        print(f"Using protobuf version {expected_version}")


def protoc_path():
    """Return path to protoc binary."""
    binary = f"protoc-{_get_protobuf_version()}" + (
        ".exe" if sys.platform == "win32" else ""
    )
    return os.path.join("bin", binary)


def extract_message_info():
    """Get information about all messages of interest."""
    filename = os.path.join(BASE_PATH, "ProtocolMessage.proto")

    with open(filename, encoding="utf-8", mode="r") as file:
        types_found = False

        for line in file:
            stripped = line.lstrip().rstrip()

            # Look for the Type enum
            if stripped == "enum Type {":
                types_found = True
                continue
            if types_found and stripped == "}":
                break
            if not types_found:
                continue

            constant = stripped.split(" ")[0]
            title = constant.title().replace("_", "").replace("Hid", "HID")  # Hack...
            accessor = title[0].lower() + title[1:]

            if not os.path.exists(os.path.join(BASE_PATH, title + ".proto")):
                continue

            yield MessageInfo(title + "_pb2", title, accessor, constant)


def extract_unreferenced_messages():
    """Get messages not referenced anywhere."""
    for filename in os.listdir(BASE_PATH):
        tmp = os.path.splitext(filename)
        if tmp[1] != ".proto" or tmp[0] == "ProtocolMessage":
            continue

        with open(os.path.join(BASE_PATH, filename), encoding="utf-8") as file:
            for line in file:
                if line.startswith("message"):
                    yield tmp[0] + "_pb2", line.split(" ")[1]


def generate_module_code():
    """Generate protobuf message wrappercode."""
    message_names = set()
    packages = []
    messages = []
    extensions = []
    constants = []

    # Extract everything needed to generate output file
    for info in extract_message_info():
        message_names.add(info.title)
        packages.append("from . import " + info.module)
        messages.append(f"from .{info.module} import {info.title}")
        extensions.append(
            f"ProtocolMessage.{info.const}: {info.module}.{info.accessor},"
        )
        constants.append(f"{info.const} = ProtocolMessage.{info.const}")

        reused = REUSED_MESSAGES.get(info.const)
        if reused: