from __future__ import annotations

from functools import lru_cache
from typing import Any

from config import OPERATION_MODE_BATCHED_STREAM, ConnectorConfig
from fivetran.operations import CheckpointTrackingOperationSink, FivetranOperationSink, OperationSink
from fivetran.schema import build_fivetran_schema
from increment.reader import sync_incremental
from increment.schema_index import StorageSchemaIndex, load_latest_storage_schema_index, load_storage_schema_index
from observability import info_with_memory
from snapshot.reader import sync_snapshot
from snapshot.tables import (
    SnapshotTables,
    load_snapshot_tables,
)
from state import ConnectorState, apply_table_resync_requests, normalize_state
from storage.object_store import ObjectStore
from storage.resolver import StorageURI, create_object_store, redact_storage_uri


@lru_cache(maxsize=8)
def _cached_store(storage_uri: str) -> ObjectStore:
    return create_object_store(StorageURI.parse(storage_uri))


def _load_common_inputs(
    configuration: dict[str, Any],
) -> tuple[ConnectorConfig, ObjectStore]:
    config = ConnectorConfig.from_dict(configuration)
    store = _cached_store(config.storage_uri)
    return config, store


def _load_schema_inputs(
    configuration: dict[str, Any],
) -> tuple[ConnectorConfig, ObjectStore, SnapshotTables, StorageSchemaIndex]:
    config, store = _load_common_inputs(configuration)
    snapshot_tables = (
        load_snapshot_tables(store, snapshot_prefix=config.snapshot_prefix)
        if config.snapshot_required
        else SnapshotTables(tables=())
    )
    schema_index = load_latest_storage_schema_index(store, config.incremental_prefix)
    return config, store, snapshot_tables, schema_index


def schema(configuration: dict[str, Any]) -> list[dict[str, object]]:
    config, _store, snapshot_tables, schema_index = _load_schema_inputs(configuration)
    info_with_memory(
        "building schema",
        phase="schema",
        storage_uri=redact_storage_uri(config.storage_uri),
        snapshot_table_count=len(snapshot_tables.tables),
        incremental_schema_count=len(schema_index.table_defs),
    )
    return build_fivetran_schema(snapshot_tables, schema_index, config)


def update(configuration: dict[str, Any], state: dict[str, Any] | str | None) -> None:
    from fivetran_connector_sdk import Operations as op

    config, store = _load_common_inputs(configuration)
    _update_loaded(
        config,
        store,
        state,
        FivetranOperationSink(op, use_batched_stream=config.operation_mode == OPERATION_MODE_BATCHED_STREAM),
    )


def _update(
    configuration: dict[str, Any],
    state: dict[str, Any] | str | None,
    operations: OperationSink,
) -> None:
    config, store = _load_common_inputs(configuration)
    _update_loaded(config, store, state, operations)


def _update_loaded(
    config: ConnectorConfig,
    store: ObjectStore,
    state: dict[str, Any] | str | None,
    operations: OperationSink,
) -> None:
    state_obj = normalize_state(
        state,
        sync_task_id=config.sync_task_id,
        snapshot_required=config.snapshot_required,
    )
    tracked_operations = CheckpointTrackingOperationSink(operations)
    info_with_memory(
        "starting update",
        phase="update",
        storage_uri=redact_storage_uri(config.storage_uri),
    )

    if state_obj.scan is not None:
        schema_index = load_storage_schema_index(store, config.incremental_prefix)
        sync_incremental(
            config=config,
            store=store,
            schema_index=schema_index,
            state=state_obj,
            operations=tracked_operations,
        )
        if state_obj.scan is not None:
            _checkpoint_if_missing(tracked_operations, state_obj)
            return

    if not config.snapshot_required:
        schema_index = load_storage_schema_index(store, config.incremental_prefix)
        sync_incremental(
            config=config,
            store=store,
            schema_index=schema_index,
            state=state_obj,
            operations=tracked_operations,
        )
        _checkpoint_if_missing(tracked_operations, state_obj)
        return

    apply_table_resync_requests(state_obj)

    sync_snapshot(
        config=config,
        store=store,
        state=state_obj,
        operations=tracked_operations,
    )

    sync_incremental(
        config=config,
        store=store,
        schema_index=load_storage_schema_index(store, config.incremental_prefix),
        state=state_obj,
        operations=tracked_operations,
    )
    _checkpoint_if_missing(tracked_operations, state_obj)


def _checkpoint_if_missing(operations: CheckpointTrackingOperationSink, state_obj: ConnectorState) -> None:
    if operations.checkpoint_called:
        return
    operations.checkpoint(state=state_obj)
    info_with_memory(
        "emitted final checkpoint because sync completed without an earlier checkpoint",
        phase="update",
    )


__all__ = ["schema", "update"]