# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from abc import ABC, abstractmethod
import logging
import subprocess
import threading
import time
import uuid

from devil.utils import reraiser_thread


class ExpensiveLineTransformer(ABC):
    def __init__(
        self, process_start_timeout, minimum_timeout, per_line_timeout
    ):
        self._process_start_timeout = process_start_timeout
        self._minimum_timeout = minimum_timeout
        self._per_line_timeout = per_line_timeout
        self._started = False
        # Allow only one thread to call TransformLines() at a time.
        self._lock = threading.Lock()
        # Ensure that only one thread attempts to kill self._proc in Close().
        self._close_lock = threading.Lock()
        self._closed_called = False
        # Assign to None so that attribute exists if Popen() throws.
        self._proc = None
        # Start process eagerly to hide start-up latency.
        self._proc_start_time = None

    def start(self):
        # delay the start of the process, to allow the initialization of the
        # descendant classes first.
        if self._started:
            logging.error(
                '%s: Trying to start an already started command', self.name
            )
            return

        # Start process eagerly to hide start-up latency.
        self._proc_start_time = time.time()

        if not self.command:
            logging.error('%s: No command available', self.name)
            return

        self._proc = subprocess.Popen(
            self.command,
            bufsize=1,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            universal_newlines=True,
            close_fds=True,
        )
        self._started = True

    def IsClosed(self):
        return (
            not self._started
            or self._closed_called
            or self._proc.returncode is not None
        )

    def IsBusy(self):
        return self._lock.locked()

    def IsReady(self):
        return self._started and not self.IsClosed() and not self.IsBusy()

    def TransformLines(self, lines):
        """Symbolizes names found in the given lines.

        If anything goes wrong (process crashes, timeout, etc), returns |lines|.

        Args:
          lines: A list of strings without trailing newlines.

        Returns:
          A list of strings without trailing newlines.
        """
        if not lines:
            return []

        # symbolized output contain more lines than the input, as the symbolized
        # stacktraces will be added. To account for the extra output lines, keep
        # reading until this eof_line token is reached. Using a format that will
        # be considered a "useful line" without modifying its output by
        # third_party/android_platform/development/scripts/stack_core.py
        eof_line = self.getEofLine()
        out_lines = []

        def _reader():
            while True:
                line = self._proc.stdout.readline()
                # Return an empty string at EOF (when stdin is closed).
                if not line:
                    break
                line = line[:-1]
                if line == eof_line:
                    break
                out_lines.append(line)

        if self.IsBusy():
            logging.warning('%s: Having to wait for transformation.', self.name)

        # Allow only one thread to operate at a time.
        with self._lock:
            if self.IsClosed():
                if self._started and not self._closed_called:
                    logging.warning(
                        '%s: Process exited with code=%d.',
                        self.name,
                        self._proc.returncode,
                    )
                    self.Close()
                return lines

            reader_thread = reraiser_thread.ReraiserThread(_reader)
            reader_thread.start()

            try:
                self._proc.stdin.write('\n'.join(lines))
                self._proc.stdin.write('\n{}\n'.format(eof_line))
                self._proc.stdin.flush()
                time_since_proc_start = time.time() - self._proc_start_time
                timeout = max(
                    0, self._process_start_timeout - time_since_proc_start
                ) + max(
                    self._minimum_timeout, len(lines) * self._per_line_timeout
                )
                reader_thread.join(timeout)
                if self.IsClosed():
                    logging.warning(
                        '%s: Close() called by another thread during join().',
                        self.name,
                    )
                    return lines
                if reader_thread.is_alive():
                    logging.error(
                        '%s: Timed out after %f seconds with input:',
                        self.name,
                        timeout,
                    )
                    for l in lines:
                        logging.error(l)
                    logging.error(eof_line)
                    logging.error('%s: End of timed out input.', self.name)
                    logging.error('%s: Timed out output was:', self.name)
                    for l in out_lines:
                        logging.error(l)
                    logging.error('%s: End of timed out output.', self.name)
                    self.Close()
                    return lines
                return out_lines
            except IOError:
                logging.exception(
                    '%s: Exception during transformation', self.name
                )
                self.Close()
                return lines

    def Close(self):
        with self._close_lock:
            needs_closing = not self.IsClosed()
            self._closed_called = True

        if needs_closing:
            self._proc.stdin.close()
            self._proc.kill()
            self._proc.wait()

    def __del__(self):
        # self._proc is None when Popen() fails.
        if not self._closed_called and self._proc:
            logging.error('%s: Forgot to Close()', self.name)
            self.Close()

    @property
    @abstractmethod
    def name(self): ...

    @property
    @abstractmethod
    def command(self): ...

    @staticmethod
    def getEofLine():
        # Use a format that will be considered a "useful line" without modifying its
        # output by third_party/android_platform/development/scripts/stack_core.py
        return "Generic useful log header: '{}'".format(uuid.uuid4().hex)


class ExpensiveLineTransformerPool(ABC):
    def __init__(self, max_restarts, pool_size, passthrough_on_failure):
        self._max_restarts = max_restarts
        self._pool = [self.CreateTransformer() for _ in range(pool_size)]
        self._passthrough_on_failure = passthrough_on_failure
        # Allow only one thread to select from the pool at a time.
        self._lock = threading.Lock()
        self._num_restarts = 0

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.Close()

    def TransformLines(self, lines):
        with self._lock:
            assert self._pool, 'TransformLines() called on a closed Pool.'

            # transformation is broken.
            if self._num_restarts == self._max_restarts:
                if self._passthrough_on_failure:
                    return lines
                raise Exception('%s is broken.' % self.name)

            # Restart any closed transformer.
            for i, d in enumerate(self._pool):
                if d.IsClosed():
                    logging.warning(
                        '%s: Restarting closed instance.', self.name
                    )
                    self._pool[i] = self.CreateTransformer()
                    self._num_restarts += 1
                    if self._num_restarts == self._max_restarts:
                        logging.warning('%s: MAX_RESTARTS reached.', self.name)
                        if self._passthrough_on_failure:
                            return lines
                        raise Exception('%s is broken.' % self.name)

            selected = next(
                (x for x in self._pool if x.IsReady()), self._pool[0]
            )
            # Rotate the order so that next caller will not choose the same one.
            self._pool.remove(selected)
            self._pool.append(selected)

        return selected.TransformLines(lines)

    def Close(self):
        with self._lock:
            for d in self._pool:
                d.Close()
            self._pool = None

    @abstractmethod
    def CreateTransformer(self): ...

    @property
    @abstractmethod
    def name(self): ...