# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.
import io
import logging
import os
import sqlite3
import sys
import threading
import time
import uuid
from dataclasses import dataclass
from functools import cached_property
from pathlib import Path

from botocore.compat import get_md5
from botocore.exceptions import MD5UnavailableError
from botocore.useragent import UserAgentComponent
from botocore.utils import ensure_boolean

from awscli.compat import is_windows
from awscli.utils import add_component_to_user_agent_extra

LOG = logging.getLogger(__name__)

_CACHE_DIR = Path.home() / '.aws' / 'cli' / 'cache'
_DATABASE_FILENAME = 'session.db'
_SESSION_LENGTH_SECONDS = 60 * 30
_SESSION_ID_LENGTH = 12
# Set to "true" to skip collecting session ids entirely. Opting out avoids
# opening the session database, whose file locking can be expensive when the
# database lives on a network filesystem.
_SESSION_ID_DISABLED_ENV_VAR = 'AWS_CLI_SESSION_ID_DISABLED'

def _get_checksum():
    hashlib_params = {"usedforsecurity": False}
    try:
        checksum = get_md5(**hashlib_params)
    except MD5UnavailableError:
        import hashlib

        checksum = hashlib.sha256(**hashlib_params)
    return checksum


@dataclass
class CLISessionData:
    key: str
    session_id: str
    timestamp: int


class CLISessionDatabaseConnection:
    _CREATE_TABLE = """
        CREATE TABLE IF NOT EXISTS session (
          key TEXT PRIMARY KEY,
          session_id TEXT NOT NULL,
          timestamp INTEGER NOT NULL
        )
    """
    _CREATE_HOST_ID_TABLE = """
        CREATE TABLE IF NOT EXISTS host_id (
          key INTEGER PRIMARY KEY,
          id TEXT UNIQUE NOT NULL
        )
    """
    _CHECK_HOST_ID = """
        SELECT COUNT(*) FROM host_id
    """
    _INSERT_HOST_ID = """
        INSERT OR IGNORE INTO host_id (
            key, id
        ) VALUES (?, ?)
    """
    _ENABLE_WAL = 'PRAGMA journal_mode=WAL'

    def __init__(self, connection=None, cache_dir=None):
        self._cache_dir = cache_dir or _CACHE_DIR
        self._ensure_cache_dir()
        self._connection = connection or sqlite3.connect(
            self._cache_dir / _DATABASE_FILENAME,
            check_same_thread=False,
            isolation_level=None,
        )
        self._ensure_database_setup()

    def execute(self, query, *parameters):
        try:
            return self._connection.execute(query, *parameters)
        except sqlite3.OperationalError:
            # Process timed out waiting for database lock.
            # Return any empty `Cursor` object instead of
            # raising an exception.
            LOG.debug(
                'Failed to execute session database query: %s',
                query,
                exc_info=True,
            )
            return sqlite3.Cursor(self._connection)

    def _ensure_cache_dir(self):
        self._cache_dir.mkdir(parents=True, exist_ok=True)

    def _ensure_database_setup(self):
        self._create_session_table()
        self._create_host_id_table()
        self._ensure_host_id()
        self._try_to_enable_wal()

    def _create_session_table(self):
        self.execute(self._CREATE_TABLE)

    def _create_host_id_table(self):
        self.execute(self._CREATE_HOST_ID_TABLE)

    def _ensure_host_id(self):
        cur = self.execute(self._CHECK_HOST_ID)
        host_id_ct = cur.fetchone()[0]
        if host_id_ct == 0:
            self.execute(
                self._INSERT_HOST_ID,
                # Hardcode `0` as primary key to ensure
                # there's only ever 1 host id in the table.
                (
                    0,
                    str(uuid.uuid4()),
                ),
            )

    def _try_to_enable_wal(self):
        try:
            self.execute(self._ENABLE_WAL)
        except sqlite3.Error:
            # This is just a performance enhancement so it is optional. Not all
            # systems will have a sqlite compiled with the WAL enabled.
            LOG.debug(
                'Failed to enable WAL for session database', exc_info=True
            )


class CLISessionDatabaseWriter:
    _WRITE_RECORD = """
        INSERT OR REPLACE INTO session (
            key, session_id, timestamp
        ) VALUES (?, ?, ?)
    """

    def __init__(self, connection):
        self._connection = connection

    def write(self, data):
        self._connection.execute(
            self._WRITE_RECORD,
            (
                data.key,
                data.session_id,
                data.timestamp,
            ),
        )


class CLISessionDatabaseReader:
    _READ_RECORD = """
        SELECT *
        FROM session
        WHERE key = ?
    """
    _READ_HOST_ID = """
        SELECT id
        FROM host_id
        WHERE key = 0
    """

    def __init__(self, connection):
        self._connection = connection

    def read(self, key):
        cursor = self._connection.execute(self._READ_RECORD, (key,))
        result = cursor.fetchone()
        if result is None:
            return
        return CLISessionData(*result)

    def read_host_id(self):
        cursor = self._connection.execute(self._READ_HOST_ID)
        return cursor.fetchone()[0]


class CLISessionDatabaseSweeper:
    _DELETE_RECORDS = """
        DELETE FROM session
        WHERE timestamp < ?
    """

    def __init__(self, connection):
        self._connection = connection

    def sweep(self, timestamp):
        try:
            self._connection.execute(self._DELETE_RECORDS, (timestamp,))
        except Exception:
            # This is just a background cleanup task. No need to
            # handle it or direct to stderr.
            LOG.debug('Failed to sweep session database', exc_info=True)
            return


class CLISessionGenerator:
    def generate_session_id(self, host_id, tty, timestamp):
        return self._generate_checksum(host_id, tty, timestamp)

    def generate_cache_key(self, host_id, tty):
        return self._generate_checksum(host_id, tty)

    def _generate_checksum(self, *args):
        checksum = _get_checksum()
        str_to_hash = ""
        for arg in args:
            if arg is not None:
                str_to_hash += str(arg)
        checksum.update(str_to_hash.encode('utf-8'))
        return checksum.hexdigest()[:_SESSION_ID_LENGTH]


class CLISessionOrchestrator:
    def __init__(self, generator, writer, reader, sweeper):
        self._generator = generator
        self._writer = writer
        self._reader = reader
        self._sweeper = sweeper

        self._sweep_cache()

    @cached_property
    def cache_key(self):
        return self._generator.generate_cache_key(self._host_id, self._tty)

    @cached_property
    def _session_id(self):
        return self._generator.generate_session_id(
            self._host_id, self._tty, self._timestamp
        )

    @cached_property
    def session_id(self):
        if (cached_data := self._reader.read(self.cache_key)) is not None:
            # Cache hit, but session id is expired. Generate new id and update.
            if (
                cached_data.timestamp + _SESSION_LENGTH_SECONDS
                < self._timestamp
            ):
                cached_data.session_id = self._session_id
            # Always update the timestamp to last used.
            cached_data.timestamp = self._timestamp
            self._writer.write(cached_data)
            return cached_data.session_id
        # Cache miss, generate and write new record.
        session_id = self._session_id
        session_data = CLISessionData(
            self.cache_key, session_id, self._timestamp
        )
        self._writer.write(session_data)
        return session_id

    @cached_property
    def _tty(self):
        # os.ttyname is only available on Unix platforms.
        if is_windows:
            return
        try:
            return os.ttyname(sys.stdin.fileno())
        except (OSError, io.UnsupportedOperation):
            # Standard input was redirected to a pseudofile.
            # This can happen when running tests on IDEs or
            # running scripts with redirected input, etc.
            return

    @cached_property
    def _host_id(self):
        return self._reader.read_host_id()

    @cached_property
    def _timestamp(self):
        return int(time.time())

    def _sweep_cache(self):
        t = threading.Thread(
            target=self._sweeper.sweep,
            args=(self._timestamp - _SESSION_LENGTH_SECONDS,),
            daemon=True,
        )
        t.start()


def _get_cli_session_orchestrator():
    conn = CLISessionDatabaseConnection()
    return CLISessionOrchestrator(
        CLISessionGenerator(),
        CLISessionDatabaseWriter(conn),
        CLISessionDatabaseReader(conn),
        CLISessionDatabaseSweeper(conn),
    )


def is_telemetry_disabled(env=None):
    if env is None:
        env = os.environ
    return ensure_boolean(env.get(_SESSION_ID_DISABLED_ENV_VAR, 'false'))


def register_session_id_event(session, orchestrator_factory=None):
    if is_telemetry_disabled():
        # Skip registering the event entirely so we never open the session
        # database. This lets users avoid the database's file locking, which
        # can be expensive on network filesystems.
        return
    if orchestrator_factory is None:
        orchestrator_factory = _get_cli_session_orchestrator
    event_emitter = session.get_component('event_emitter')

    def _inject_session_id(**kwargs):
        try:
            orchestrator = orchestrator_factory()
            sid_component = UserAgentComponent(
                "sid", orchestrator.session_id
            ).to_string()
            # Insert sid after md/installer to preserve original
            # user-agent component ordering.
            extra = session.user_agent_extra
            idx = extra.find('md/installer')
            end = extra.find(' ', idx)
            if end == -1:
                end = len(extra)
            session.user_agent_extra = (
                extra[:end] + f' {sid_component}' + extra[end:]
            )
        except Exception:
            # Ideally, the AWS CLI should never throw if the session id
            # can't be generated since it's not critical for users. Issues
            # with session data should instead be caught server-side.
            LOG.debug('Failed to generate session id', exc_info=True)
        event_emitter.unregister('before-create-client', _inject_session_id)

    event_emitter.register('before-create-client', _inject_session_id)
