import os

import pytest

from test.integration.connectors.utils.constants import SOURCE_TAG, UNCATEGORIZED_TAG
from test.integration.connectors.utils.validation.source import (
    SourceValidationConfigs,
    get_all_file_data,
    run_all_validations,
    update_fixtures,
)
from unstructured_ingest.v2.interfaces import Downloader, Indexer
from unstructured_ingest.v2.processes.connectors.notion.connector import (
    CONNECTOR_TYPE,
    NotionAccessConfig,
    NotionConnectionConfig,
    NotionDownloader,
    NotionDownloaderConfig,
    NotionIndexer,
    NotionIndexerConfig,
)


@pytest.mark.tags(SOURCE_TAG, CONNECTOR_TYPE, UNCATEGORIZED_TAG)
def test_notion_source_database(temp_dir):
    # Retrieve environment variables
    notion_api_key = os.environ["NOTION_API_KEY"]

    # Create connection and indexer configurations
    access_config = NotionAccessConfig(notion_api_key=notion_api_key)
    connection_config = NotionConnectionConfig(
        access_config=access_config,
    )
    index_config = NotionIndexerConfig(
        database_ids=["1722c3765a0a8082b382ebc2c62d3f4c"], recursive=False
    )

    download_config = NotionDownloaderConfig(download_dir=temp_dir)

    # Instantiate indexer and downloader
    indexer = NotionIndexer(
        connection_config=connection_config,
        index_config=index_config,
    )
    downloader = NotionDownloader(
        connection_config=connection_config,
        download_config=download_config,
    )

    # Run the source connector validation
    source_connector_validation(
        indexer=indexer,
        downloader=downloader,
        configs=SourceValidationConfigs(
            test_id="notion_database",
            expected_num_files=1,
            validate_downloaded_files=True,
            exclude_fields_extend=["metadata.date_created", "metadata.date_modified"],
        ),
    )


@pytest.mark.tags(SOURCE_TAG, CONNECTOR_TYPE, UNCATEGORIZED_TAG)
def test_notion_source_page(temp_dir):
    # Retrieve environment variables
    notion_api_key = os.environ["NOTION_API_KEY"]

    # Create connection and indexer configurations
    access_config = NotionAccessConfig(notion_api_key=notion_api_key)
    connection_config = NotionConnectionConfig(
        access_config=access_config,
    )
    index_config = NotionIndexerConfig(
        page_ids=["1572c3765a0a806299f0dd6999f9e4c7"], recursive=False
    )

    download_config = NotionDownloaderConfig(download_dir=temp_dir)

    # Instantiate indexer and downloader
    indexer = NotionIndexer(
        connection_config=connection_config,
        index_config=index_config,
    )
    downloader = NotionDownloader(
        connection_config=connection_config,
        download_config=download_config,
    )

    # Run the source connector validation
    source_connector_validation(
        indexer=indexer,
        downloader=downloader,
        configs=SourceValidationConfigs(
            test_id="notion_page",
            expected_num_files=1,
            validate_downloaded_files=True,
            exclude_fields_extend=["metadata.date_created", "metadata.date_modified"],
        ),
    )


@pytest.mark.tags(SOURCE_TAG, CONNECTOR_TYPE, UNCATEGORIZED_TAG)
def source_connector_validation(
    indexer: Indexer,
    downloader: Downloader,
    configs: SourceValidationConfigs,
    overwrite_fixtures: bool = os.getenv("OVERWRITE_FIXTURES", "False").lower() == "true",
) -> None:
    # Run common validations on the process of running a source connector, supporting dynamic
    # validators that get passed in along with comparisons on the saved expected values.
    # If overwrite_fixtures is st to True, will ignore all validators but instead overwrite the
    # expected values with what gets generated by this test.
    all_predownload_file_data = []
    all_postdownload_file_data = []
    indexer.precheck()
    download_dir = downloader.download_config.download_dir
    test_output_dir = configs.test_output_dir()

    for file_data in indexer.run():
        assert file_data
        predownload_file_data = file_data.model_copy(deep=True)
        all_predownload_file_data.append(predownload_file_data)
        resp = downloader.run(file_data=file_data)
        if isinstance(resp, list):
            for r in resp:
                postdownload_file_data = r["file_data"].model_copy(deep=True)
                all_postdownload_file_data.append(postdownload_file_data)
        else:
            postdownload_file_data = resp["file_data"].model_copy(deep=True)
            all_postdownload_file_data.append(postdownload_file_data)

    if not overwrite_fixtures:
        print("Running validation")
        run_all_validations(
            configs=configs,
            predownload_file_data=all_predownload_file_data,
            postdownload_file_data=all_postdownload_file_data,
            download_dir=download_dir,
            test_output_dir=test_output_dir,
        )
    else:
        print("Running fixtures update")
        update_fixtures(
            output_dir=test_output_dir,
            download_dir=download_dir,
            all_file_data=get_all_file_data(
                all_predownload_file_data=all_predownload_file_data,
                all_postdownload_file_data=all_postdownload_file_data,
            ),
            save_downloads=configs.validate_downloaded_files,
            save_filedata=configs.validate_file_data,
        )
