#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#
from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING, List

from dagger import Container, Directory

from pipelines.airbyte_ci.connectors.build_image.steps.python_connectors import BuildConnectorImages
from pipelines.airbyte_ci.connectors.consts import CONNECTOR_TEST_STEP_ID
from pipelines.airbyte_ci.connectors.context import ConnectorContext, PipelineContext
from pipelines.airbyte_ci.connectors.reports import Report
from pipelines.consts import LOCAL_BUILD_PLATFORM
from pipelines.dagger.actions.python.poetry import with_poetry
from pipelines.helpers.connectors.command import run_connector_steps
from pipelines.helpers.execution.run_steps import STEP_TREE, StepToRun
from pipelines.models.secrets import Secret
from pipelines.models.steps import Step, StepResult, StepStatus

if TYPE_CHECKING:
    from anyio import Semaphore


def _get_erd_folder(code_directory: Path) -> Path:
    path = code_directory / "erd"
    if not path.exists():
        path.mkdir()
    return path


class GenerateDbml(Step):
    context: ConnectorContext

    title = "Generate DBML file"

    def __init__(self, context: PipelineContext, skip_relationship_generation: bool) -> None:
        super().__init__(context)
        self._skip_relationship_generation = skip_relationship_generation

    async def _run(self, connector_to_discover: Container) -> StepResult:
        if not self._skip_relationship_generation and not self.context.genai_api_key:
            raise ValueError("GENAI_API_KEY needs to be provided if the relationship generation is not skipped")

        discovered_catalog = await self._get_discovered_catalog(connector_to_discover)

        connector_directory = await self.context.get_connector_dir()
        command = ["poetry", "run", "erd", "--source-path", "/source", "--source-technical-name", self.context.connector.technical_name]
        if self._skip_relationship_generation:
            command.append("--skip-llm-relationships")

        erd_directory = self._build_erd_container(connector_directory, discovered_catalog).with_exec(command).directory("/source/erd")
        await erd_directory.export(str(_get_erd_folder(self.context.connector.code_directory)))

        return StepResult(step=self, status=StepStatus.SUCCESS, output=erd_directory)

    async def _get_discovered_catalog(self, connector_to_discover: Container) -> str:
        source_config_path_in_container = "/data/config.json"
        discover_output = (
            await connector_to_discover.with_new_file(source_config_path_in_container, contents=self._get_default_config().value)
            .with_exec(["discover", "--config", source_config_path_in_container], use_entrypoint=True)
            .stdout()
        )
        return self._get_schema_from_discover_output(discover_output)

    def _get_default_config(self) -> Secret:
        if not self.context.local_secret_store:
            raise ValueError("Expecting local secret store to be set up but couldn't find it")

        filtered_secrets = [secret for secret in self.context.local_secret_store.get_all_secrets() if secret.file_name == "config.json"]
        if not filtered_secrets:
            raise ValueError("Expecting at least one secret to match name `config.json`")
        elif len(filtered_secrets) > 1:
            self.logger.warning(
                f"Expecting only one secret with name `config.json but got {len(filtered_secrets)}. Will take the first one in the list."
            )

        return filtered_secrets[0]

    @staticmethod
    def _get_schema_from_discover_output(discover_output: str) -> str:
        for line in discover_output.split("\n"):
            json_line = json.loads(line)
            if json_line.get("type") == "CATALOG":
                return json.dumps(json.loads(line).get("catalog"))
        raise ValueError("No catalog was found in output")

    def _build_erd_container(self, connector_directory: Directory, discovered_catalog: str) -> Container:
        """Create a container to run ERD generation."""
        container = with_poetry(self.context)
        if self.context.genai_api_key:
            container = container.with_secret_variable("GENAI_API_KEY", self.context.genai_api_key.as_dagger_secret(self.dagger_client))

        container = (
            container.with_mounted_directory("/source", connector_directory)
            .with_new_file("/source/erd/discovered_catalog.json", contents=discovered_catalog)
            .with_mounted_directory("/app", self.context.erd_package_dir)
            .with_workdir("/app")
        )

        return container.with_exec(["poetry", "lock"], use_entrypoint=True).with_exec(["poetry", "install"], use_entrypoint=True)


class UploadDbmlSchema(Step):
    context: ConnectorContext

    title = "Upload DBML file to dbdocs.io"

    def __init__(self, context: PipelineContext) -> None:
        super().__init__(context)

    async def _run(self, erd_directory: Directory) -> StepResult:
        if not self.context.dbdocs_token:
            raise ValueError(
                "In order to publish to dbdocs, DBDOCS_TOKEN needs to be provided. Either pass the value or skip the publish step"
            )

        dbdocs_container = (
            self.dagger_client.container()
            .from_("node:lts-bullseye-slim")
            .with_exec(["npm", "install", "-g", "dbdocs"], use_entrypoint=True)
            .with_env_variable("DBDOCS_TOKEN", self.context.dbdocs_token.value)
            .with_workdir("/airbyte_dbdocs")
            .with_file("/airbyte_dbdocs/dbdocs.dbml", erd_directory.file("source.dbml"))
        )

        db_docs_build = ["dbdocs", "build", "dbdocs.dbml", f"--project={self.context.connector.technical_name}"]
        await dbdocs_container.with_exec(db_docs_build).stdout()
        # TODO: produce link to dbdocs in output logs

        return StepResult(step=self, status=StepStatus.SUCCESS)


async def run_connector_generate_erd_pipeline(
    context: ConnectorContext,
    semaphore: "Semaphore",
    skip_steps: List[str],
) -> Report:
    context.targeted_platforms = [LOCAL_BUILD_PLATFORM]
    steps_to_run: STEP_TREE = []

    steps_to_run.append([StepToRun(id=CONNECTOR_TEST_STEP_ID.BUILD, step=BuildConnectorImages(context))])

    steps_to_run.append(
        [
            StepToRun(
                id=CONNECTOR_TEST_STEP_ID.DBML_FILE,
                step=GenerateDbml(context, CONNECTOR_TEST_STEP_ID.LLM_RELATIONSHIPS in skip_steps),
                args=lambda results: {"connector_to_discover": results[CONNECTOR_TEST_STEP_ID.BUILD].output[LOCAL_BUILD_PLATFORM]},
                depends_on=[CONNECTOR_TEST_STEP_ID.BUILD],
            ),
        ]
    )

    if CONNECTOR_TEST_STEP_ID.PUBLISH_ERD not in skip_steps:
        steps_to_run.append(
            [
                StepToRun(
                    id=CONNECTOR_TEST_STEP_ID.PUBLISH_ERD,
                    step=UploadDbmlSchema(context),
                    args=lambda results: {"erd_directory": results[CONNECTOR_TEST_STEP_ID.DBML_FILE].output},
                    depends_on=[CONNECTOR_TEST_STEP_ID.DBML_FILE],
                ),
            ]
        )

    return await run_connector_steps(context, semaphore, steps_to_run)
