import {
    type AthenaClient,
    CreateWorkGroupCommand,
    DeleteWorkGroupCommand,
    GetQueryExecutionCommand,
    GetQueryResultsCommand,
    type GetQueryResultsCommandOutput,
    GetWorkGroupCommand,
    ListDatabasesCommand,
    ListQueryExecutionsCommand,
    ListTableMetadataCommand,
    type ListTableMetadataCommandOutput,
    ListWorkGroupsCommand,
    type QueryExecution,
    StartQueryExecutionCommand,
    StopQueryExecutionCommand,
    type TableMetadata,
    type WorkGroup,
    type WorkGroupSummary,
} from '@aws-sdk/client-athena'
import {NotFoundError, RuntimeError, ValidationError} from '../cloud-spi/errors'
import {
    awsQuerySchema,
    PRIMARY_WORKGROUP,
    QUERY_HISTORY_LIMIT,
    QUERY_RESULT_ROW_LIMIT,
    WORKGROUP_NAME_PATTERN,
} from '../cloud-spi/querySchema'
import type {
    CloudResource,
    CloudServiceAdapter,
    CreateResourceInput,
    QueryCatalog,
    QueryExecutionSummary,
    QueryInput,
    QueryResult,
    ResourceQuery,
    ServiceSchema,
} from '../cloud-spi/types'

/** Athena's own catalog name for the Glue Data Catalog. */
const DATA_CATALOG = 'AwsDataCatalog'
/** Athena's cap on a query string. */
const MAX_QUERY_LENGTH = 262_144
const S3_URI = /^s3:\/\/[a-z0-9.-]{3,63}(\/.*)?$/

/**
 * Talks to Amazon Athena, verified against Floci core (community, `latest`).
 *
 * Notes from probing the runtime:
 *  - Queries really run: TheCloud executes the SQL with DuckDB over the files a
 *    Glue table points at, and writes the result to S3 as CSV, as Athena does.
 *  - Statements that return no rows, such as SHOW TABLES, fail in the runtime,
 *    so the editor lists databases and tables through the catalog APIs instead.
 *  - Data scanned is always reported as 0 bytes, so the editor shows only the
 *    run time.
 */
export class AwsQueryAdapter implements CloudServiceAdapter {
    readonly cloud = 'aws' as const
    readonly service = 'query' as const

    constructor(private readonly athena: AthenaClient) {}

    schema(): ServiceSchema {
        return awsQuerySchema()
    }

    async list(query: ResourceQuery = {}): Promise<CloudResource[]> {
        const found: WorkGroupSummary[] = []
        let next: string | undefined
        do {
            const res = await this.athena.send(new ListWorkGroupsCommand({MaxResults: 50, ...(next ? {NextToken: next} : {})}))
            found.push(...(res.WorkGroups ?? []))
            next = res.NextToken
        } while (next)

        const search = query.search?.trim().toLowerCase()
        return found.map(toResource).filter((workGroup) => !search
            || workGroup.name.toLowerCase().includes(search)
            || String(workGroup.metadata.description ?? '').toLowerCase().includes(search))
    }

    async get(id: string): Promise<CloudResource | null> {
        try {
            const workGroup = (await this.athena.send(new GetWorkGroupCommand({WorkGroup: id}))).WorkGroup
            return workGroup ? toResource(workGroup) : null
        } catch (error) {
            if (errorName(error) === 'InvalidRequestException') return null
            throw error
        }
    }

    async create(input: CreateResourceInput): Promise<CloudResource> {
        const values = input.values
        const name = text(values.name)
        if (!new RegExp(WORKGROUP_NAME_PATTERN).test(name)) throw new ValidationError('name must be up to 128 letters, numbers, dots, dashes or underscores')
        if (name === PRIMARY_WORKGROUP) throw new ValidationError('Every account already has the primary workgroup')
        const outputLocation = text(values.outputLocation)
        if (outputLocation && !S3_URI.test(outputLocation)) throw new ValidationError('outputLocation must be an S3 URI such as s3://my-bucket/athena-results/')
        const description = text(values.description)

        await this.athena.send(new CreateWorkGroupCommand({
            Name: name,
            ...(description ? {Description: description} : {}),
            Configuration: {
                ...(outputLocation ? {ResultConfiguration: {OutputLocation: outputLocation.endsWith('/') ? outputLocation : `${outputLocation}/`}} : {}),
                EnforceWorkGroupConfiguration: values.enforce === 'true',
            },
        }))
        const created = await this.get(name)
        if (!created) throw new RuntimeError(`Athena did not return the new workgroup ${name}`)
        return created
    }

    async delete(id: string): Promise<void> {
        if (id === PRIMARY_WORKGROUP) throw new ValidationError('The primary workgroup cannot be deleted')
        await this.athena.send(new DeleteWorkGroupCommand({WorkGroup: id, RecursiveDeleteOption: true}))
    }

    async startQuery(workGroup: string, input: QueryInput): Promise<QueryExecutionSummary> {
        const sql = typeof input.sql === 'string' ? input.sql.trim() : ''
        if (!sql) throw new ValidationError('Enter a query to run')
        if (sql.length > MAX_QUERY_LENGTH) throw new ValidationError(`A query can be at most ${MAX_QUERY_LENGTH} characters`)
        const database = text(input.database)

        const started = await this.athena.send(new StartQueryExecutionCommand({
            QueryString: sql,
            WorkGroup: workGroup,
            ...(database ? {QueryExecutionContext: {Catalog: DATA_CATALOG, Database: database}} : {}),
        }))
        if (!started.QueryExecutionId) throw new RuntimeError('Athena did not return a query id')
        return (await this.execution(workGroup, started.QueryExecutionId))
    }

    async getQueryResult(workGroup: string, queryId: string): Promise<QueryResult> {
        const execution = await this.execution(workGroup, queryId)
        if (execution.state !== 'SUCCEEDED' || !execution.outputLocation) return {execution, columns: [], rows: [], truncated: false}

        const pages: GetQueryResultsCommandOutput[] = []
        let next: string | undefined
        let rowCount = 0
        do {
            const page = await this.athena.send(new GetQueryResultsCommand({QueryExecutionId: queryId, MaxResults: 1000, ...(next ? {NextToken: next} : {})}))
            pages.push(page)
            rowCount += page.ResultSet?.Rows?.length ?? 0
            next = page.NextToken
        } while (next && rowCount <= QUERY_RESULT_ROW_LIMIT)

        const columnInfo = pages[0]?.ResultSet?.ResultSetMetadata?.ColumnInfo ?? []
        const raw = pages.flatMap((page) => page.ResultSet?.Rows ?? []).map((row) => (row.Data ?? []).map((cell) => cell.VarCharValue ?? null))
        const header = raw[0] ?? []
        const columns = columnInfo.length > 0
            ? columnInfo.map((column) => ({name: column.Label ?? column.Name ?? '', type: column.Type ?? null}))
            : header.map((name) => ({name: name ?? '', type: null}))
        // Athena repeats the column names as the first row of a SELECT's results.
        const body = raw.length > 0 && header.length === columns.length && header.every((name, index) => name === columns[index].name) ? raw.slice(1) : raw
        return {execution, columns, rows: body.slice(0, QUERY_RESULT_ROW_LIMIT), truncated: body.length > QUERY_RESULT_ROW_LIMIT || Boolean(next)}
    }

    async stopQuery(workGroup: string, queryId: string): Promise<void> {
        await this.execution(workGroup, queryId)
        await this.athena.send(new StopQueryExecutionCommand({QueryExecutionId: queryId}))
    }

    async listQueries(workGroup: string): Promise<QueryExecutionSummary[]> {
        const ids = (await this.athena.send(new ListQueryExecutionsCommand({WorkGroup: workGroup, MaxResults: QUERY_HISTORY_LIMIT}))).QueryExecutionIds ?? []
        const executions = await Promise.all(ids.slice(0, QUERY_HISTORY_LIMIT).map(async (id) => {
            const res = await this.athena.send(new GetQueryExecutionCommand({QueryExecutionId: id}))
            return res.QueryExecution ? toExecution(res.QueryExecution) : null
        }))
        return executions
            .filter((execution): execution is QueryExecutionSummary => execution !== null && (execution.workGroup ?? PRIMARY_WORKGROUP) === workGroup)
            .sort((a, b) => (b.submittedAt ?? '').localeCompare(a.submittedAt ?? ''))
    }

    async describeQueryCatalog(database?: string): Promise<QueryCatalog> {
        const databases: string[] = []
        let next: string | undefined
        do {
            const res = await this.athena.send(new ListDatabasesCommand({CatalogName: DATA_CATALOG, MaxResults: 50, ...(next ? {NextToken: next} : {})}))
            databases.push(...(res.DatabaseList ?? []).map((entry) => entry.Name ?? '').filter(Boolean))
            next = res.NextToken
        } while (next)
        databases.sort()

        const chosen = database && databases.includes(database) ? database : databases[0] ?? null
        if (!chosen) return {databases, database: null, tables: []}

        const tables: TableMetadata[] = []
        let tableToken: string | undefined
        do {
            const res: ListTableMetadataCommandOutput = await this.athena.send(new ListTableMetadataCommand({CatalogName: DATA_CATALOG, DatabaseName: chosen, MaxResults: 50, ...(tableToken ? {NextToken: tableToken} : {})}))
            tables.push(...(res.TableMetadataList ?? []))
            tableToken = res.NextToken
        } while (tableToken)
        return {
            databases,
            database: chosen,
            tables: tables.map((table) => ({
                name: table.Name ?? '',
                columns: [...(table.Columns ?? []), ...(table.PartitionKeys ?? [])].map((column) => ({name: column.Name ?? '', type: column.Type ?? ''})),
            })).sort((a, b) => a.name.localeCompare(b.name)),
        }
    }

    /** A query, refused unless it ran in this workgroup, so one workgroup's page never shows another's. */
    private async execution(workGroup: string, queryId: string): Promise<QueryExecutionSummary> {
        let execution: QueryExecution | undefined
        try {
            execution = (await this.athena.send(new GetQueryExecutionCommand({QueryExecutionId: queryId}))).QueryExecution
        } catch (error) {
            if (errorName(error) === 'InvalidRequestException') throw new NotFoundError(`Query ${queryId} not found`)
            throw error
        }
        if (!execution || (execution.WorkGroup ?? PRIMARY_WORKGROUP) !== workGroup) throw new NotFoundError(`Query ${queryId} is not in the ${workGroup} workgroup`)
        return toExecution(execution)
    }
}

function toResource(workGroup: WorkGroup | WorkGroupSummary): CloudResource {
    const name = workGroup.Name ?? ''
    const detail = 'Configuration' in workGroup ? workGroup : undefined
    const configuration = detail?.Configuration
    return {
        id: name,
        name,
        cloud: 'aws',
        service: 'query',
        type: 'athenaWorkGroup',
        region: null,
        createdAt: workGroup.CreationTime ? workGroup.CreationTime.toISOString() : null,
        status: workGroup.State ?? null,
        metadata: {
            description: workGroup.Description || null,
            ...(detail ? {
                outputLocation: configuration?.ResultConfiguration?.OutputLocation ?? null,
                enforceConfiguration: configuration?.EnforceWorkGroupConfiguration === true,
                engineVersion: configuration?.EngineVersion?.EffectiveEngineVersion ?? configuration?.EngineVersion?.SelectedEngineVersion ?? null,
            } : {}),
        },
    }
}

function toExecution(execution: QueryExecution): QueryExecutionSummary {
    const status = execution.Status
    return {
        id: execution.QueryExecutionId ?? '',
        sql: execution.Query ?? '',
        database: execution.QueryExecutionContext?.Database ?? null,
        workGroup: execution.WorkGroup ?? null,
        state: status?.State ?? 'QUEUED',
        stateReason: status?.StateChangeReason ?? status?.AthenaError?.ErrorMessage ?? null,
        statementType: execution.StatementType ?? null,
        submittedAt: status?.SubmissionDateTime ? status.SubmissionDateTime.toISOString() : null,
        completedAt: status?.CompletionDateTime ? status.CompletionDateTime.toISOString() : null,
        runtimeMs: execution.Statistics?.TotalExecutionTimeInMillis ?? execution.Statistics?.EngineExecutionTimeInMillis ?? null,
        dataScannedBytes: execution.Statistics?.DataScannedInBytes ?? null,
        outputLocation: execution.ResultConfiguration?.OutputLocation ?? null,
    }
}

function text(value: unknown): string {
    return typeof value === 'string' ? value.trim() : ''
}

function errorName(error: unknown): string | undefined {
    return typeof error === 'object' && error !== null ? (error as {name?: string}).name : undefined
}