# Copyright (c) 2022 Alex Butler, Daniel Purtov
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this
# software and associated documentation files (the "Software"), to deal in the Software
# without restriction, including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons
# to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or
# substantial portions of the Software.
"""Airbyte tap class"""

from __future__ import annotations

import errno
import os
import subprocess
import sys
import time
from contextlib import contextmanager
from datetime import date, datetime
from decimal import Decimal
from enum import Enum
from functools import lru_cache
from pathlib import Path, PurePath
from queue import Empty, Queue
from tempfile import TemporaryDirectory
from threading import Lock, Thread
from typing import Any, Callable, Dict, Iterable, List, Optional
from uuid import UUID

import click
import orjson
import singer_sdk._singerlib as singer
from singer_sdk import Stream, Tap
from singer_sdk import typing as th
from singer_sdk.cli import common_options
from singer_sdk.helpers._classproperty import classproperty

# Sentinel value for broken pipe
PIPE_CLOSED = object()


def default(obj):
    if isinstance(obj, (datetime, date)):
        return obj.isoformat()
    elif isinstance(obj, Decimal):
        return float(obj)
    elif isinstance(obj, UUID):
        return str(obj)
    elif isinstance(obj, bytes):
        return obj.decode("utf-8")
    elif isinstance(obj, Enum):
        return obj.value
    return str(obj)


def write_message(message) -> None:
    try:
        sys.stdout.buffer.write(
            orjson.dumps(message.to_dict(), option=TapAirbyte.ORJSON_OPTS, default=default)
        )
        sys.stdout.buffer.flush()
    except IOError as ex:
        # Broken pipe
        if ex.errno == errno.EPIPE and TapAirbyte.pipe_status is not PIPE_CLOSED:
            TapAirbyte.logger.info("Received SIGPIPE, stopping sync of stream.")
            TapAirbyte.pipe_status = PIPE_CLOSED
            # Prevent BrokenPipe writes to closed stdout
            os.dup2(os.open(os.devnull, os.O_WRONLY), sys.stdout.fileno())
        else:
            raise


STDOUT_LOCK = Lock()
singer.write_message = write_message


class AirbyteException(Exception):
    pass


class AirbyteMessage(str, Enum):
    RECORD = "RECORD"
    STATE = "STATE"
    LOG = "LOG"
    TRACE = "TRACE"
    CATALOG = "CATALOG"
    SPEC = "SPEC"
    CONNECTION_STATUS = "CONNECTION_STATUS"


# These translate between Singer's replication method and Airbyte's sync mode
REPLICATION_METHOD_MAP = {
    "FULL_TABLE": "FULL_REFRESH",
    "INCREMENTAL": "INCREMENTAL",
    "LOG_BASED": "INCREMENTAL",
}
# We are piping to Singer targets, so this field is irrelevant
NOOP_AIRBYTE_SYNC_MODE = "append"


class TapAirbyte(Tap):
    name = "tap-airbyte"
    config_jsonschema = th.PropertiesList(
        th.Property(
            "airbyte_config",
            th.ObjectType(),
            required=False,
            description=(
                "Configuration to pass through to the Airbyte tap, this can be gleaned"
                " by running the the tap with the `--about` flag and the `--config` flag pointing"
                " to a file containing the `airbyte_spec` configuration. This is a JSON object."
            ),
        ),
        th.Property(
            "file_mounts",
            th.ArrayType(
                th.ObjectType(
                    th.Property(
                        "source",
                        th.StringType,
                        required=True,
                        description="Source path to mount",
                    ),
                    th.Property(
                        "target",
                        th.StringType,
                        required=True,
                        description="Target path to mount",
                    ),
                )
            ),
            required=False,
            description=(
                "File mounts to make available to the Airbyte tap. Expects a list of maps"
                " containing source and target."
            ),
        ),
    ).to_dict()
    conf_dir: str = "/tmp"
    pipe_status = None

    _file_mounts: Optional[List[Dict[str, str]]] = None

    # Airbyte -> Demultiplexer -< Singer Streams
    singer_consumers: List[Thread] = []
    buffers: Dict[str, Queue] = {}

    # State container
    airbyte_state: Dict[str, Any] = {}

    ORJSON_OPTS = orjson.OPT_APPEND_NEWLINE

    @classproperty
    def cli(cls) -> Callable:
        @common_options.PLUGIN_VERSION
        @common_options.PLUGIN_ABOUT
        @common_options.PLUGIN_CONFIG
        @click.option(
            "--format",
            "-f",
            help="Specify output style for --about [text|json|markdown]",
        )
        @click.option(
            "--discover",
            "-d",
            is_flag=True,
            help="Run the tap in discovery mode.",
        )
        @click.option(
            "--test",
            "-t",
            is_flag=True,
            help="Use --test to run the Airbyte connection test.",
        )
        @click.option(
            "--catalog",
            "-c",
            help="Use a Singer catalog file with the tap.",
            type=click.Path(),
        )
        @click.option(
            "--state",
            "-s",
            help="Use a bookmarks file for incremental replication.",
            type=click.Path(),
        )
        @click.command(
            help="Execute the Singer tap.",
            context_settings={"help_option_names": ["--help"]},
        )
        def cli(
            version: bool = False,
            about: bool = False,
            discover: bool = False,
            test: bool = False,
            config: tuple[str, ...] = (),
            state: Optional[str] = None,
            catalog: Optional[str] = None,
            format: Optional[str] = None,
        ) -> None:
            if version:
                cls.print_version()
                return
            if not about:
                cls.print_version(print_fn=cls.logger.info)
            validate_config: bool = True
            if discover or about:
                validate_config = False
            parse_env_config = False
            config_files: list[PurePath] = []
            for config_path in config:
                if config_path == "ENV":
                    parse_env_config = True
                    continue
                if not Path(config_path).is_file():
                    raise FileNotFoundError(
                        f"Could not locate config file at '{config_path}'."
                        "Please check that the file exists."
                    )
                config_files.append(Path(config_path))
            # Enrich about info with spec if possible
            if about:
                cls.discover_streams = lambda _: []
                try:
                    # pylint: disable=not-callable
                    tap: TapAirbyte = cls(  # type: ignore
                        config=config_files or None,
                        state=state,
                        catalog=catalog,
                        parse_env_config=parse_env_config,
                        validate_config=validate_config,
                    )
                    # pylint: enable=not-callable
                    spec = tap.run_spec()["connectionSpecification"]
                except Exception:
                    cls.logger.info("Tap-Airbyte instantiation failed. Printing basic about info.")
                    cls.print_about(output_format=format)
                else:
                    cls.logger.info(
                        "Tap-Airbyte instantiation succeeded. Printing spec-enriched about info."
                    )
                    cls.config_jsonschema["properties"]["airbyte_config"] = spec
                    cls.print_about(output_format=format)
                    cls.print_spec_as_config(spec)
                return
            # End modification
            # pylint: disable=not-callable
            tap: TapAirbyte = cls(  # type: ignore
                config=config_files or None,
                state=state,
                catalog=catalog,
                parse_env_config=parse_env_config,
                validate_config=validate_config,
            )
            # pylint: enable=not-callable
            if discover:
                tap.run_discovery()
                if test:
                    tap.run_connection_test()
            elif test:
                tap.run_connection_test()
            else:
                tap.sync_all()

        return cli

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    def run_help(self):
        subprocess.run(
            ["/airbyte/base.sh", "--help"],
            check=True,
        )

    def run_spec(self):
        proc = subprocess.run(
            ["/airbyte/base.sh", "spec"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True
        )
        for line in proc.stdout.decode("utf-8").splitlines():
            try:
                message = orjson.loads(line)
            except orjson.JSONDecodeError:
                if line:
                    self.logger.warning("Could not parse message: %s", line)
                continue
            if message["type"] in (AirbyteMessage.LOG, AirbyteMessage.TRACE):
                self._process_log_message(message)
            elif message["type"] == AirbyteMessage.SPEC:
                return message["spec"]
            else:
                self.logger.warning("Unhandled message: %s", message)
        if proc.returncode != 0:
            raise AirbyteException("Could not run spec for tap")
        raise AirbyteException(
            "Could not output spec, no spec message received.\n"
            f"Stdout: {proc.stdout.decode('utf-8')}\n"
            f"Stderr: {proc.stderr.decode('utf-8')}"
        )

    @staticmethod
    def print_spec_as_config(spec: Dict[str, Any]) -> None:
        print("\nSetup Instructions:\n")
        print("airbyte_config:")
        for prop, schema in spec["properties"].items():
            if "description" in schema:
                print(f"  # {schema['description']}")
            print(f"  {prop}: {'fixme' if schema['type'] != 'object' else ''}")
            if schema["type"] == "object":
                if "oneOf" in schema:
                    for i, one_of in enumerate(schema["oneOf"]):
                        print(f"    # Option {i + 1}")
                        for inner_prop, inner_schema in one_of["properties"].items():
                            if inner_prop == "option_title":
                                continue
                            if "description" in inner_schema:
                                print(f"    # {inner_schema['description']}")
                            print(f"    {inner_prop}: fixme")
                else:
                    for inner_prop, inner_schema in schema["properties"].items():
                        if "description" in inner_schema:
                            print(f"    # {inner_schema['description']}")
                        print(f"    {inner_prop}: fixme")

    @staticmethod
    def mount_files(mounts: List[Dict[str, str]]) -> str:
        if len(mounts) > 0:
            for file in mounts:
                subprocess.run(
                    [
                        "cp",
                        "-r",
                        f"{file['source']}",
                        f"{file['target']}",
                    ],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    check=True,
                )

    def run_check(self) -> bool:
        with TemporaryDirectory() as tmpdir:
            with open(f"{tmpdir}/config.json", "wb") as f:
                f.write(orjson.dumps(self.config.get("airbyte_config", {})))
            self.mount_files(
                [{"source": f"{tmpdir}/config.json", "target": f"{self.conf_dir}/config.json"}]
                + self.config.get("file_mounts", [])
            )
            proc = subprocess.run(
                [
                    "./base.sh",
                    "check",
                    "--config",
                    f"{self.conf_dir}/config.json",
                ],
                cwd="/airbyte",
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
            )
        for line in proc.stdout.decode("utf-8").splitlines():
            try:
                message = orjson.loads(line)
            except orjson.JSONDecodeError:
                if line:
                    self.logger.warning("Could not parse message: %s", line)
                continue
            if message["type"] in (AirbyteMessage.LOG, AirbyteMessage.TRACE):
                self._process_log_message(message)
            elif message["type"] == AirbyteMessage.CONNECTION_STATUS:
                if message["connectionStatus"]["status"] == "SUCCEEDED":
                    self.logger.info(
                        "Configuration has been verified via the Airbyte check command."
                    )
                    return True
                else:
                    self.logger.error(
                        "Connection check failed: %s",
                        message["connectionStatus"]["message"],
                    )
                    return False
            else:
                self.logger.warning("Unhandled message: %s", message)
        if proc.returncode != 0:
            raise AirbyteException(
                f"Connection check failed with return code {proc.returncode}:"
                f" {proc.stderr.decode()}"
            )
        raise AirbyteException(
            "Could not verify connection, no connection status message received.\n"
            f"Stdout: {proc.stdout.decode('utf-8')}\n"
            f"Stderr: {proc.stderr.decode('utf-8')}"
        )

    def run_connection_test(self) -> bool:  # type: ignore
        return self.run_check()

    @contextmanager
    def run_read(self):
        with TemporaryDirectory() as tmpdir:
            with (
                open(f"{tmpdir}/config.json", "wb") as config,
                open(f"{tmpdir}/catalog.json", "wb") as catalog,
            ):
                config.write(orjson.dumps(self.config.get("airbyte_config", {})))
                catalog.write(orjson.dumps(self.configured_airbyte_catalog))
            if self.airbyte_state:
                with open(f"{tmpdir}/state.json", "wb") as state:
                    self.logger.debug("Using state: %s", self.airbyte_state)
                    state.write(orjson.dumps(self.airbyte_state))

            mounts = (
                [
                    {"source": f"{tmpdir}/config.json", "target": f"{self.conf_dir}/config.json"},
                    {"source": f"{tmpdir}/catalog.json", "target": f"{self.conf_dir}/catalog.json"},
                ]
                + self.config.get("file_mounts", [])
                + (
                    {"source": f"{tmpdir}/state.json", "target": f"{self.conf_dir}/state.json"}
                    if self.airbyte_state
                    else []
                )
            )
            self.mount_files(mounts)

            proc = subprocess.Popen(
                [
                    "./base.sh",
                    "read",
                    "--config",
                    f"{self.conf_dir}/config.json",
                    "--catalog",
                    f"{self.conf_dir}/catalog.json",
                ]
                + (["--state", f"{self.conf_dir}/state.json"] if self.airbyte_state else []),
                cwd="/airbyte",
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            try:
                # Context is held until EOF or exception
                yield proc
            finally:
                if not self.eof_received:
                    proc.kill()
                    self.logger.warning("Airbyte process terminated before EOF message received.")
                self.logger.debug("Waiting for Airbyte process to terminate.")
                returncode = proc.wait()
                if not self.eof_received and TapAirbyte.pipe_status is not PIPE_CLOSED:
                    # If EOF was not received, the process was killed and we should raise an exception
                    type, value, _ = sys.exc_info()
                    raise AirbyteException(
                        f"Airbyte process terminated early:\n{type.__name__}: {value}"
                    )
                if returncode != 0 and TapAirbyte.pipe_status is not PIPE_CLOSED:
                    # If EOF was received, the process should have exited with return code 0
                    raise AirbyteException(
                        f"Airbyte process failed with return code {returncode}:"
                        f" {proc.stderr.read()}"
                    )

    def _process_log_message(self, airbyte_message: Dict[str, Any]) -> None:
        if airbyte_message["type"] == AirbyteMessage.LOG:
            self.logger.info(airbyte_message["log"])
        elif airbyte_message["type"] == AirbyteMessage.TRACE:
            if airbyte_message["trace"].get("type") == "ERROR":
                exc = AirbyteException(
                    airbyte_message["trace"]["error"].get("stack_trace", "Airbyte process failed.")
                )
                self.logger.critical(
                    airbyte_message["trace"]["error"]["message"],
                    exc_info=exc,
                )
                raise exc
            self.logger.debug(airbyte_message["trace"])

    @property
    @lru_cache
    def airbyte_catalog(self):
        with TemporaryDirectory() as tmpdir:
            with open(f"{tmpdir}/config.json", "wb") as f:
                f.write(orjson.dumps(self.config.get("airbyte_config", {})))
            self.mount_files(
                [{"source": f"{tmpdir}/config.json", "target": f"{self.conf_dir}/config.json"}]
                + self.config.get("file_mounts", [])
            )
            proc = subprocess.run(
                [
                    "./base.sh",
                    "discover",
                    "--config",
                    f"{self.conf_dir}/config.json",
                ],
                cwd="/airbyte",
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True
            )
        for line in proc.stdout.decode("utf-8").splitlines():
            try:
                message = orjson.loads(line)
            except orjson.JSONDecodeError:
                continue
            if message["type"] in (AirbyteMessage.LOG, AirbyteMessage.TRACE):
                self._process_log_message(message)
            elif message["type"] == AirbyteMessage.CATALOG:
                return message["catalog"]
        if proc.returncode != 0:
            raise AirbyteException(
                f"Discover failed with return code {proc.returncode}: {proc.stderr.decode()}"
            )
        raise AirbyteException(
            "Could not discover catalog, no catalog message received. \n"
            f"Stdout: {proc.stdout.decode('utf-8')}\n"
            f"Stderr: {proc.stderr.decode('utf-8')}"
        )

    @property
    def configured_airbyte_catalog(self) -> dict:
        output = {"streams": []}
        for stream in self.airbyte_catalog["streams"]:
            entry = self.catalog.get_stream(stream["name"])
            if entry is None:
                continue
            if entry.metadata.root.selected is False:
                continue
            try:
                sync_mode = REPLICATION_METHOD_MAP.get(
                    entry.replication_method.upper(), stream["supported_sync_modes"][0]
                )
                if sync_mode.lower() not in stream["supported_sync_modes"]:
                    sync_mode = stream["supported_sync_modes"][0]
            except (IndexError, KeyError):
                sync_mode = "FULL_REFRESH"
            output["streams"].append(
                {
                    "stream": stream,
                    "sync_mode": sync_mode.lower(),
                    "destination_sync_mode": NOOP_AIRBYTE_SYNC_MODE,
                }
            )
        return output

    def load_state(self, state: Dict[str, Any]) -> None:
        super().load_state(state)
        self.airbyte_state = state

    def sync_all(self) -> None:  # type: ignore
        stream: Stream
        self.eof_received = False
        for stream in self.streams.values():
            if not stream.selected and not stream.has_selected_descendents:
                self.logger.info("Skipping deselected stream '%s'.", stream.name)
                continue
            consumer = Thread(target=stream.sync, daemon=True)
            consumer.start()
            self.singer_consumers.append(consumer)
        t1 = time.perf_counter()
        with self.run_read() as airbyte_job:
            # Main processor loop
            while TapAirbyte.pipe_status is not PIPE_CLOSED:
                message = airbyte_job.stdout.readline()
                if not message and airbyte_job.poll() is not None:
                    self.eof_received = True
                    break
                try:
                    airbyte_message = orjson.loads(message)
                except orjson.JSONDecodeError:
                    if message:
                        self.logger.warning("Could not parse message: %s", message)
                    continue
                if airbyte_message["type"] == AirbyteMessage.RECORD:
                    stream_buffer: Queue = self.buffers.setdefault(
                        airbyte_message["record"]["stream"],
                        Queue(),
                    )
                    stream_buffer.put_nowait(airbyte_message["record"]["data"])
                elif airbyte_message["type"] in (
                    AirbyteMessage.LOG,
                    AirbyteMessage.TRACE,
                ):
                    self._process_log_message(airbyte_message)
                elif airbyte_message["type"] == AirbyteMessage.STATE:
                    state_message = airbyte_message["state"]
                    if "data" in state_message:
                        unpacked_state = state_message["data"]
                    elif "type" == "STREAM":
                        unpacked_state = state_message["stream"]
                    elif "type" == "GLOBAL":
                        unpacked_state = state_message["global"]
                    elif "type" == "LEGACY":
                        unpacked_state = state_message["legacy"]
                    self.airbyte_state = unpacked_state
                    with STDOUT_LOCK:
                        singer.write_message(singer.StateMessage(self.airbyte_state))
                else:
                    self.logger.warning("Unhandled message: %s", airbyte_message)
        # Daemon threads will be terminated when the main thread exits
        # so we do not need to wait on them to join after SIGPIPE
        if TapAirbyte.pipe_status is not PIPE_CLOSED:
            self.logger.info("Waiting for sync threads to finish...")
            for sync in self.singer_consumers:
                sync.join()
            # Write final state if EOF was received from Airbyte
            if self.eof_received:
                with STDOUT_LOCK:
                    singer.write_message(singer.StateMessage(self.airbyte_state))
        t2 = time.perf_counter()
        for stream in self.streams.values():
            stream.log_sync_costs()
        self.logger.info("Synced %s streams in %0.2f seconds.", len(self.streams), t2 - t1)

    def discover_streams(self) -> List[Stream]:
        output_streams: List[AirbyteStream] = []
        stream: Dict[str, Any]
        for stream in self.airbyte_catalog["streams"]:
            airbyte_stream = AirbyteStream(
                tap=self,
                name=stream["name"],
                schema=stream["json_schema"],
            )
            try:
                # this is [str, ...?] in the Airbyte catalog
                if "cursor_field" in stream and isinstance(stream["cursor_field"][0], str):
                    airbyte_stream.replication_key = stream["cursor_field"][0]
            except IndexError:
                pass
            try:
                # this is [[str, ...]] in the Airbyte catalog
                if "primary_key" in stream and isinstance(stream["primary_key"][0], List):
                    airbyte_stream.primary_keys = stream["primary_key"][0]
                elif "source_defined_primary_key" in stream and isinstance(
                    stream["source_defined_primary_key"][0], List
                ):
                    airbyte_stream.primary_keys = stream["source_defined_primary_key"][0]
            except IndexError:
                pass
            output_streams.append(airbyte_stream)
        return output_streams


class AirbyteStream(Stream):
    """Stream class for Airbyte streams."""

    def __init__(self, tap: TapAirbyte, schema: dict, name: str) -> None:
        super().__init__(tap, schema, name)
        self.parent = tap
        self._buffer: Optional[Queue] = None

    def _write_record_message(self, record: dict) -> None:
        for record_message in self._generate_record_messages(record):
            with STDOUT_LOCK:
                singer.write_message(record_message)

    def _write_state_message(self) -> None:
        pass

    def _increment_stream_state(self, *args, **kwargs) -> None:
        pass

    @property
    def buffer(self) -> Queue:
        """Get the buffer for the stream."""
        if not self._buffer:
            while self.name not in self.parent.buffers:
                if self.parent.eof_received:
                    # EOF received, no records for this stream
                    self._buffer = Queue()
                    break
                self.logger.debug(f"Waiting for records from Airbyte for stream {self.name}...")
                time.sleep(1)
            else:
                self._buffer = self.parent.buffers[self.name]
        return self._buffer

    def get_records(self, context: Optional[dict]) -> Iterable[dict]:
        """Get records from the stream."""
        while (
            self.parent.eof_received is False or not self.buffer.empty()
        ) and TapAirbyte.pipe_status is not PIPE_CLOSED:
            try:
                # The timeout permits the consumer to re-check the producer is alive
                yield self.buffer.get(timeout=1.0)
            except Empty:
                continue
            self.buffer.task_done()
        if self.name in self.parent.buffers:
            while not self.buffer.empty() and TapAirbyte.pipe_status is not PIPE_CLOSED:
                try:
                    yield self.buffer.get(timeout=1.0)
                except Empty:
                    break
                self.buffer.task_done()


if __name__ == "__main__":
    TapAirbyte.cli()
