#!/usr/bin/python3

"""Command line tool for the Google Nest Smart Device Management API

You must configure your device as described:
  https://developers.google.com/nest/device-access/get-started

which will give you a project_id, client_id, and client_secret.  This tool
will do a one time setup to get an access token, then from there on will
cache the token on local disk.

Once authenticated, you can run commands like:

$ google_nest --project_id=<project_id> list
$ google_nest --project_id=<project_id> get <device_id>
"""

import argparse
import asyncio
import errno
import json
import logging
import os
import pickle
import socket

import yaml
from aiohttp import ClientSession, TCPConnector
from google.auth.credentials import Credentials
from google.auth.transport.requests import Request
from google.oauth2 import service_account
from google_auth_oauthlib.flow import InstalledAppFlow

from .auth import AbstractAuth
from .camera_traits import CameraLiveStreamTrait
from .event import EventMessage
from .google_nest_api import GoogleNestAPI
from .google_nest_subscriber import AsyncEventCallback, GoogleNestSubscriber
from .thermostat_traits import (
    ThermostatEcoTrait,
    ThermostatModeTrait,
    ThermostatTemperatureSetpointTrait,
)

# Define command line arguments
parser = argparse.ArgumentParser(
    description="Command line tool for Google Nest SDM API"
)
parser.add_argument("--project_id", required=True, help="Device Access program id")
parser.add_argument("--creds_json_file", help="Service acccount credentials file")
parser.add_argument("--client_id", help="OAuth credentials client_id")
parser.add_argument("--client_secret", help="OAuth credentials client_secret")
parser.add_argument(
    "--token_cache",
    help="File storage for long lived creds",
    default="~/.config/google_nest/token_cache",
)
parser.add_argument(
    "-v", "--verbose", help="Increase output verbosity", action="store_true"
)
parser.add_argument(
    "--output_type",
    type=str,
    choices=["json", "yaml"],
    help="Change the output type from json or yaml (default).",
    default="yaml",
)

cmd_parser = parser.add_subparsers(dest="command", required=True)
list_structures_parser = cmd_parser.add_parser("list_structures")
list_devices_parser = cmd_parser.add_parser("list_devices")
get_structure_parser = cmd_parser.add_parser("get_structure")
get_structure_parser.add_argument("structure_id")
get_device_parser = cmd_parser.add_parser("get_device")
get_device_parser.add_argument("device_id")
set_mode_parser = cmd_parser.add_parser(
    "set_mode", description="Change the thermostat mode."
)
set_mode_parser.add_argument("device_id")
set_mode_parser.add_argument(
    "mode",
    help="The mode to change the thermostat to.",
    choices=["MANUAL_ECO", "HEAT", "COOL", "HEATCOOL", "OFF"],
)
set_heat_parser = cmd_parser.add_parser(
    "set_heat", description="Sets the target temperature when in HEAT mode."
)
set_heat_parser.add_argument("device_id")
set_heat_parser.add_argument("heat", type=float)
set_cool_parser = cmd_parser.add_parser(
    "set_cool", help="Sets the target temperature when in COOL mode."
)
set_cool_parser.add_argument("device_id")
set_cool_parser.add_argument(
    "cool",
    type=float,
    help="The target temperature to set when the thermostat is in COOL mode.",
)
set_range_parser = cmd_parser.add_parser(
    "set_range", help="Sets the min/max temperature when in HEATCOOL mode."
)
set_range_parser.add_argument("device_id")
set_range_parser.add_argument(
    "heat", type=float, help="The minimum target temperature to set."
)
set_range_parser.add_argument(
    "cool", type=float, help="The maximum target temperature to set."
)
generate_rtsp_stream_parser = cmd_parser.add_parser("generate_rtsp_stream")
generate_rtsp_stream_parser.add_argument("device_id")
subscribe_parser = cmd_parser.add_parser("subscribe")
subscribe_parser.add_argument("subscription_id")
subscribe_parser.add_argument("device_id", nargs="?")

OAUTH2_AUTHORIZE = (
    "https://nestservices.google.com/partnerconnections/{project_id}/auth"
)
OAUTH2_TOKEN = "https://www.googleapis.com/oauth2/v4/token"
SDM_SCOPES = [
    "https://www.googleapis.com/auth/sdm.service",
    "https://www.googleapis.com/auth/pubsub",
]
API_URL = "https://smartdevicemanagement.googleapis.com/v1"


class Auth(AbstractAuth):
    """Implementation of AbstractAuth that uses the token cache."""

    def __init__(
        self,
        websession: ClientSession,
        user_creds: Credentials,
        service_creds: Credentials,
        api_url: str,
    ):
        """Initialize Google Nest Device Access auth."""
        super().__init__(websession, api_url)
        self._user_creds = user_creds
        self._service_creds = service_creds

    async def async_get_access_token(self):
        """Return a valid access token."""
        return self._user_creds.token

    async def async_get_creds(self):
        return self._service_creds


def CreateCreds(args) -> Credentials:
    """Runs an interactive flow to get OAuth creds."""

    creds = None
    token_cache = os.path.expanduser(args.token_cache)
    if os.path.exists(token_cache):
        with open(token_cache, "rb") as token:
            creds = pickle.load(token)

    # If there are no (valid) credentials available, let the user log in.
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            if not args.client_id or not args.client_secret:
                raise ValueError("Required flag --client_id or --client_secret missing")
            client_config = {
                "installed": {
                    "client_id": args.client_id,
                    "client_secret": args.client_secret,
                    "redirect_uris": ["urn:ietf:wg:oauth:2.0:oob"],
                    "auth_uri": OAUTH2_AUTHORIZE.format(project_id=args.project_id),
                    "token_uri": OAUTH2_TOKEN,
                },
            }
            app_flow = InstalledAppFlow.from_client_config(
                client_config, scopes=SDM_SCOPES
            )
            creds = app_flow.run_console()
        # Save the credentials for the next run
        if not os.path.exists(os.path.dirname(token_cache)):
            try:
                os.makedirs(os.path.dirname(token_cache))
            except OSError as exc:  # Guard against race condition
                if exc.errno != errno.EEXIST:
                    raise
        with open(token_cache, "wb") as token:
            pickle.dump(creds, token)
    return creds


def PrintStructure(structure, output_type):
    if output_type == "json":
        print(json.dumps(structure.raw_data))
    else:
        print(yaml.dump(structure.raw_data))


def PrintDevice(device, output_type):
    if output_type == "json":
        print(json.dumps(device.raw_data))
    else:
        print(yaml.dump(device.raw_data))


class SubscribeCallback(AsyncEventCallback):
    def __init__(self, output_type=None):
        self._output_type = output_type

    async def async_handle_event(self, event_message: EventMessage):
        if self._output_type == "json":
            print(json.dumps(event_message.raw_data))
        else:
            print(yaml.dump(event_message.raw_data))


class DeviceWatcherCallback(AsyncEventCallback):
    def __init__(self, device, output_type):
        self._device = device
        self._output_type = output_type

    async def async_handle_event(self, event_message: EventMessage):
        print(f"event_id: {event_message.event_id}")
        print("Current device state:")
        PrintDevice(self._device, self._output_type)
        print("")


async def RunTool(args, user_creds: Credentials, service_creds: Credentials):
    conn = TCPConnector(
        family=socket.AF_INET,
        verify_ssl=False,
    )
    async with ClientSession(connector=conn) as client:
        auth = Auth(client, user_creds, service_creds, API_URL)
        api = GoogleNestAPI(auth, args.project_id)

        if args.command == "list_structures":
            structures = await api.async_get_structures()
            for structure in structures:
                PrintStructure(structure, args.output_type)
            return

        if args.command == "get_structure":
            structure = await api.async_get_structure(args.structure_id)
            PrintStructure(structure, args.output_type)
            return

        if args.command == "list_devices":
            devices = await api.async_get_devices()
            for device in devices:
                PrintDevice(device, args.output_type)
            return

        if args.command == "subscribe":
            logging.info("Subscription: %s", args.subscription_id)
            subscriber = GoogleNestSubscriber(
                auth, args.project_id, args.subscription_id
            )
            if args.device_id:
                device_manager = await subscriber.async_get_device_manager()
                device = device_manager.devices[args.device_id]
                callback = DeviceWatcherCallback(device, args.output_type)
                device.add_event_callback(callback)
            else:
                subscriber.set_update_callback(SubscribeCallback(args.output_type))
            await subscriber.start_async()
            try:
                while True:
                    await asyncio.sleep(10)
            except KeyboardInterrupt:
                subscriber.stop_async()

        # All other commands require a device_id
        device = await api.async_get_device(args.device_id)

        if args.command == "get_device":
            PrintDevice(device, args.output_type)

        if args.command == "set_mode":
            mode = args.mode
            trait = device.traits[ThermostatModeTrait.NAME]
            if mode == "MANUAL_ECO":
                trait = device.traits[ThermostatEcoTrait.NAME]
            resp = await trait.set_mode(mode)
            print(await resp.text())

        if args.command == "set_heat":
            trait = device.traits[ThermostatTemperatureSetpointTrait.NAME]
            resp = await trait.set_heat(args.heat)
            print(await resp.text())

        if args.command == "set_cool":
            trait = device.traits[ThermostatTemperatureSetpointTrait.NAME]
            resp = await trait.set_cool(args.cool)
            print(await resp.text())

        if args.command == "set_range":
            trait = device.traits[ThermostatTemperatureSetpointTrait.NAME]
            resp = await trait.set_range(args.heat, args.cool)
            print(await resp.text())

        if args.command == "generate_rtsp_stream":
            trait = device.traits[CameraLiveStreamTrait.NAME]
            stream = await trait.generate_rtsp_stream()
            print(f"URL: {stream.rtsp_stream_url}")
            print(f"Stream Token: {stream.stream_token}")
            print(f"Expires At: {stream.expires_at}")


def main():
    args = parser.parse_args()
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    user_creds = CreateCreds(args)
    if args.creds_json_file:
        service_account.Credentials.from_service_account_file(
            args.creds_json_file
        ).with_scopes(scopes=SDM_SCOPES)
    else:
        service_creds = user_creds
    loop = asyncio.get_event_loop()
    loop.run_until_complete(RunTool(args, user_creds, service_creds))
    loop.close()


if __name__ == "__main__":
    main()
