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) -> 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 provider_names:
            context = self.session_factory.resolve_context(profile=profile, region=region)
            for table_name in provider_names:
                provider = self.registry.get(table_name)
                assert provider is not None
                self.backend.ensure_provider(provider)
                if refresh or not self.backend.is_fresh(
                    provider,
                    profile=context.profile,
                    source_region=context.region,
                    ttl_seconds=self.ttl_seconds,
                ):
                    rows = provider.fetch(context)
                    self.backend.replace_rows(
                        provider,
                        profile=context.profile,
                        source_region=context.region,
                        rows=rows,
                    )

        return self.backend.execute(sql)
