# Copyright 2015 Internap.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from _socket import timeout, gaierror
import logging
import time

import paramiko
from netman.adapters import shell

from netman.adapters.shell.base import TerminalClient
from netman.core.objects.exceptions import CouldNotConnect, ConnectTimeout, CommandTimeout


class SshClient(TerminalClient):

    def __init__(self, host, username, password, port=22, prompt=('>', '#'), connect_timeout=None, command_timeout=None,
                 reading_interval=0.01, reading_chunk_size=9999):
        self.logger = logging.getLogger(__name__)

        self.host = host
        self.port = port
        self.username = username
        self.prompt = prompt
        self.command_timeout = command_timeout or shell.default_command_timeout
        connect_timeout = connect_timeout or shell.default_connect_timeout
        self.reading_interval = reading_interval
        self.reading_chunk_size = reading_chunk_size

        self.current_buffer = ''
        self.client = None
        self.channel = None
        self.full_log = ""

        self._open_channel(host, port, username, password, connect_timeout)

    def do(self, command, wait_for=None, include_last_line=False):
        self.logger.debug("[SSH][{}@{}:{}] Send >> {}".format(self.username, self.host, self.port, command))

        self.channel.send(command + '\n')
        return self._read_until(wait_for, include_last_line)

    def send_key(self, key, wait_for=None, include_last_line=False):
        self.logger.debug("[SSH][{}@{}:{}] Send KEY >> {}".format(self.username, self.host, self.port, key))

        self.channel.send(key)
        return self._read_until(wait_for, include_last_line)

    def quit(self, command):
        self.logger.debug("[SSH][{}@{}:{}] Quit >> {}".format(self.username, self.host, self.port, command))

        self.channel.send(command + '\n')

    def get_current_prompt(self):
        return self.current_buffer.splitlines()[-1]

    def _open_channel(self, host, port, username, password, connect_timeout):
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            self.client.connect(host,
                                port=port,
                                username=username,
                                password=password,
                                timeout=connect_timeout,
                                allow_agent=False,
                                look_for_keys=False)
        except timeout:
            raise ConnectTimeout(host, port)
        except gaierror:
            raise CouldNotConnect(host, port)

        self.channel = self.client.invoke_shell()

        self._wait_for(self.prompt)

    def _read_until(self, wait_for, include_last_line):
        self._wait_for(wait_for or self.prompt)

        lines = self.current_buffer.splitlines()[1:]
        if not include_last_line:
            lines = lines[:-1]

        return filter(None, lines)

    def _wait_for(self, wait_for):
        self.current_buffer = ''

        started_at = time.time()
        while not self.current_buffer.endswith(wait_for):
            while not self.channel.recv_ready():
                if time.time() - started_at > self.command_timeout:
                    raise CommandTimeout(wait_for, self.current_buffer)
                time.sleep(self.reading_interval)

            read = self.channel.recv(self.reading_chunk_size)
            self.logger.debug("[SSH][{}@{}:{}] Recv << {}".format(self.username, self.host, self.port, repr(read)))
            self.full_log += read
            self.current_buffer += read
