# 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.

"""Persistent storage of test results."""
from __future__ import print_function

import datetime
import io
import os.path
import re
import subprocess
import sys

import sqlalchemy
from sqlalchemy import orm
import subunit.v2
from subunit2sql.db import api as db_api
from subunit2sql import read_subunit
from subunit2sql import shell
from subunit2sql import write_subunit
import testtools

from stestr.repository import abstract as repository
from stestr import utils


def atomicish_rename(source, target):
    if os.name != "posix" and os.path.exists(target):
        os.remove(target)
    os.rename(source, target)


class RepositoryFactory(repository.AbstractRepositoryFactory):

    def initialise(klass, url):
        """Create a repository at url/path."""
        print("WARNING: The SQL repository type is still experimental. You "
              "might encounter issues while using it.", file=sys.stderr)
        result = Repository(url)
        # TODO(mtreinish): Figure out the python api to run the migrations for
        # setting up the schema.
        proc = subprocess.Popen(['subunit2sql-db-manage',
                                 '--database-connection', url, 'upgrade',
                                 'head'],
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
        out, err = proc.communicate()
        sys.stdout.write(str(out))
        sys.stderr.write(str(err))

        return result

    def open(self, url):
        repo = Repository(url)
        # To test the repository's existence call get_ids_for_all_tests()
        # if it raises an OperationalError that means the DB doesn't exist or
        # it couldn't connect, either way the repository was not found.
        try:
            session = repo.session_factory()
            db_api.get_ids_for_all_tests(session=session)
            session.close()
        except sqlalchemy.exc.OperationalError:
            raise repository.RepositoryNotFound(url)
        return repo


class Repository(repository.AbstractRepository):
    """subunit2sql based storage of test results.

    This repository stores each stream in a subunit2sql DB. Refer to the
    subunit2sql documentation for
    """

    def __init__(self, url):
        """Create a subunit2sql-based repository object for the repo at 'url'.

        :param base: The path to the repository.
        """
        self.base = url
        self.engine = sqlalchemy.create_engine(url)
        self.session_factory = orm.sessionmaker(bind=self.engine)

    # TODO(mtreinish): We need to add a subunit2sql api to get the count
    def count(self):
        super(Repository, self).count()

    def _get_latest_run(self):
        session = self.session_factory()
        latest_run = db_api.get_latest_run(session)
        session.close()
        if not latest_run:
            raise KeyError("No tests in repository")
        return latest_run

    def latest_id(self):
        return self._get_latest_run().uuid

    def get_failing(self):
        latest_run = self._get_latest_run()
        session = self.session_factory()
        failed_test_runs = db_api.get_test_runs_by_status_for_run_ids(
            'fail', [latest_run.id], session=session)
        session.close()
        return _Subunit2SqlRun(self.base, None, test_runs=failed_test_runs)

    def get_test_run(self, run_id):
        return _Subunit2SqlRun(self.base, run_id)

    def _get_inserter(self, partial, run_id=None, metadata=None):
        return _SqlInserter(self, partial, run_id, metadata)

    def _get_test_times(self, test_ids):
        result = {}
        # TODO(mtreinish): after subunit2sql adds a bulk query for getting
        # multiple tests by test_id at once remove the for loop
        session = self.session_factory()
        for test_id in test_ids:
            stripped_test_id = utils.cleanup_test_name(test_id)
            test = db_api.get_test_by_test_id(stripped_test_id,
                                              session=session)
            if test:
                # NOTE(mtreinish): We need to make sure the test_id with attrs
                # is used in the output dict, otherwise the scheduler won't
                # see it
                result[test_id] = test.run_time
        session.close()
        return result

    def find_metadata(self, metadata):
        session = self.session_factory()
        runs = db_api.get_runs_by_key_value('stestr_run_meta', metadata,
                                            session=session)
        return [x.uuid for x in runs]


class _Subunit2SqlRun(repository.AbstractTestRun):
    """A test run that was inserted into the repository."""

    def __init__(self, url, run_id, test_runs=None):
        engine = sqlalchemy.create_engine(url)
        self.session_factory = orm.sessionmaker(bind=engine)
        self._run_id = run_id
        self._test_runs = test_runs

    def get_id(self):
        return self._run_id

    def get_subunit_stream(self):
        stream = io.BytesIO()
        if self._run_id:
            session = self.session_factory()
            test_runs = db_api.get_tests_run_dicts_from_run_id(self._run_id,
                                                               session)
            session.close()
        else:
            test_runs = self._test_runs
        output = subunit.v2.StreamResultToBytes(stream)
        output.startTestRun()
        for test_id in test_runs:
            test = test_runs[test_id]
            # NOTE(mtreinish): test_run_metadata is not guaranteed to be
            # present for the test run
            metadata = test.get('metadata', None)
            write_subunit.write_test(output, test['start_time'],
                                     test['stop_time'], test['status'],
                                     test_id, metadata)
        output.stopTestRun()
        stream.seek(0)
        return stream

    def get_test(self):
        stream = self.get_subunit_stream()
        case = subunit.ByteStreamToStreamResult(stream)
        return case

    def get_metadata(self):
        if self._run_id:
            session = self.session_factory()
            metadata = db_api.get_run_metadata(self._run_id, session=session)
            for meta in metadata:
                if meta.key == 'stestr_run_meta':
                    return meta.value
        return None


class _SqlInserter(repository.AbstractTestRun):
    """Insert test results into a sql repository."""

    def __init__(self, repository, partial=False, run_id=None, metadata=None):
        self._repository = repository
        self.partial = partial
        self._subunit = None
        self._run_id = run_id
        self._metadata = metadata
        # Create a new session factory
        self.engine = sqlalchemy.create_engine(self._repository.base)
        self.session_factory = orm.sessionmaker(bind=self.engine,
                                                autocommit=True)

    def startTestRun(self):
        self._subunit = io.BytesIO()
        self.subunit_stream = subunit.v2.StreamResultToBytes(self._subunit)
        self.hook = testtools.CopyStreamResult([
            testtools.StreamToDict(self._handle_test),
            self.subunit_stream])
        self.hook.startTestRun()
        self.start_time = datetime.datetime.utcnow()
        session = self.session_factory()
        if not self._run_id:
            self.run = db_api.create_run(session=session)
            if self._metadata:
                db_api.add_run_metadata({'stestr_run_meta': self._metadata},
                                        self.run.id, session=session)
            self._run_id = self.run.uuid
        else:
            int_id = db_api.get_run_id_from_uuid(self._run_id, session=session)
            self.run = db_api.get_run_by_id(int_id, session=session)
        session.close()
        self.totals = {}

    def _update_test(self, test_dict, session, start_time, stop_time):
        test_id = utils.cleanup_test_name(test_dict['id'])
        db_test = db_api.get_test_by_test_id(test_id, session)
        if not db_test:
            if test_dict['status'] == 'success':
                success = 1
                fails = 0
            elif test_dict['status'] == 'fail':
                fails = 1
                success = 0
            else:
                fails = 0
                success = 0
            run_time = read_subunit.get_duration(start_time, stop_time)
            db_test = db_api.create_test(test_id, (success + fails), success,
                                         fails, run_time,
                                         session)
        else:
            test_dict['start_time'] = start_time
            test_dict['end_time'] = stop_time
            test_values = shell.increment_counts(db_test, test_dict)
            # If skipped nothing to update
            if test_values:
                db_api.update_test(test_values, db_test.id, session)
        return db_test

    def _get_attrs(self, test_id):
        attr_regex = re.compile(r'\[(.*)\]')
        matches = attr_regex.search(test_id)
        attrs = None
        if matches:
            attrs = matches.group(1)
        return attrs

    def _handle_test(self, test_dict):
        start, end = test_dict.pop('timestamps')
        if test_dict['status'] == 'exists' or None in (start, end):
            return
        elif test_dict['id'] == 'process-returncode':
            return
        session = self.session_factory()
        try:
            # Update the run counts
            if test_dict['status'] not in self.totals:
                self.totals[test_dict['status']] = 1
            else:
                self.totals[test_dict['status']] += 1
            values = {}
            if test_dict['status'] in ('success', 'xfail'):
                values['passes'] = self.totals['success']
            elif test_dict['status'] in ('fail', 'uxsuccess'):
                values['fails'] = self.totals['fail']
            elif test_dict['status'] == 'skip':
                values['skips'] = self.totals['skip']
            db_api.update_run(values, self.run.id, session=session)
            # Update the test totals
            db_test = self._update_test(test_dict, session, start,
                                        end)
            # Add the test run
            test_run = db_api.create_test_run(db_test.id, self.run.id,
                                              test_dict['status'],
                                              start, end,
                                              session)
            metadata = {}
            attrs = self._get_attrs(test_dict['id'])
            if attrs:
                metadata['attrs'] = attrs
            if test_dict.get('tags', None):
                metadata['tags'] = ",".join(test_dict['tags'])
            if metadata:
                db_api.add_test_run_metadata(
                    metadata, test_run.id, session)
            # TODO(mtreinish): Add attachments support to the DB.
            session.close()
        except Exception:
            session.rollback()
            raise

    def stopTestRun(self):
        self.hook.stopTestRun()
        stop_time = datetime.datetime.utcnow()
        self._subunit.seek(0)
        values = {}
        values['run_time'] = read_subunit.get_duration(self.start_time,
                                                       stop_time)
        session = self.session_factory()
        db_api.update_run(values, self.run.id, session=session)
        session.close()

    def status(self, *args, **kwargs):
        self.hook.status(*args, **kwargs)

    def get_id(self):
        return self._run_id
