from __future__ import annotations

from dataclasses import dataclass

from awspipe.aws import DoctorStatus, SessionFactory
from awspipe.backend import BackendAdapter, QueryResult
from awspipe.exceptions import AwspipeError, QueryPreparationError
from awspipe.providers.aws_tables import StaticProviderRegistry
from awspipe.providers.base import TableProvider
from awspipe.sql import extract_table_names


@dataclass(frozen=True)
class InspectRow:
    column: str
    type_name: str
    description: str


class QueryEngine:
    def __init__(
        self,
        backend: BackendAdapter,
        session_factory: SessionFactory | None = None,
        registry: StaticProviderRegistry | None = None,
        ttl_seconds: int = 300,
    ) -> None:
        self.backend = backend
        self.session_factory = session_factory or SessionFactory()
        self.registry = registry or StaticProviderRegistry()
        self.ttl_seconds = ttl_seconds

    def available_tables(self) -> list[TableProvider]:
        return sorted(self.registry.all(), key=lambda provider: provider.name)

    def inspect_table(self, table_name: str) -> tuple[TableProvider, list[InspectRow]]:
        provider = self.registry.get(table_name)
        if provider is None:
            raise AwspipeError(f"Unknown table: {table_name}")
        return provider, [
            InspectRow(column=column.name, type_name=column.duckdb_type, description=column.description)
            for column in provider.columns
        ]

    def doctor(self, profile: str | None, region: str | None) -> DoctorStatus:
        return self.session_factory.doctor(profile=profile, region=region, cache_path=str(self.backend.db_path))

    def query(
        self,
        sql: str,
        profile: str | None,
        region: str | None,
        refresh: bool = False,
        all_regions: bool = False,
    ) -> QueryResult:
        try:
            referenced_tables = extract_table_names(sql)
        except ValueError as exc:
            raise QueryPreparationError(str(exc)) from exc

        provider_names = [table_name for table_name in referenced_tables if self.registry.get(table_name)]
        if not provider_names:
            return self.backend.execute(sql)

        context = self.session_factory.resolve_context(profile=profile, region=region)
        regional_source_regions = self.session_factory.resolve_regions(
            context.session,
            region=context.region,
            all_regions=all_regions,
        )
        scoped_source_regions = [context.region] if not all_regions else regional_source_regions

        for table_name in provider_names:
            provider = self.registry.get(table_name)
            assert provider is not None
            self.backend.ensure_provider(provider)
            target_regions = regional_source_regions if provider.scope == "regional" else [context.region]
            for source_region in target_regions:
                provider_context = context if source_region == context.region else self.session_factory.resolve_context(
                    profile=profile,
                    region=source_region,
                )
                if refresh or not self.backend.is_fresh(
                    provider,
                    profile=provider_context.profile,
                    source_region=source_region,
                    ttl_seconds=self.ttl_seconds,
                ):
                    rows = provider.fetch(provider_context)
                    self.backend.replace_rows(
                        provider,
                        profile=provider_context.profile,
                        source_region=source_region,
                        rows=rows,
                    )

        return self.backend.execute_scoped(
            sql,
            profile=context.profile,
            table_names=provider_names,
            source_regions=scoped_source_regions,
        )
