//
//  DuckDBStore.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

import DuckDB
import Foundation

/// DuckDB-backed data store (CLAUDE.md §6). File loading goes through
/// DuckDB readers (`read_parquet`, `read_csv`) and query results are
/// materialized into `ZQDataFrame` columns: every numeric database type
/// widens to Float64, text stays as strings.
public actor ZQDataStore {
    private let database: Database
    private let connection: Connection

    public init() throws {
        self.database = try Database(store: .inMemory)
        self.connection = try database.connect()
    }

    // MARK: - Loading

    /// Loads a dataset file into a data frame. Format is inferred from the
    /// file extension: .parquet, .csv, .tsv, .json, .arrow.
    public func load(contentsOf url: URL) throws -> ZQDataFrame {
        let path = url.path
        guard FileManager.default.fileExists(atPath: path) else {
            throw ZQDataError("file not found: \(path)")
        }
        // Stata files use the native reader (CLAUDE.md §6).
        if ["dta"].contains(url.pathExtension.lowercased()) {
            return try DTAReader.read(contentsOf: url)
        }

        let escaped = path.replacingOccurrences(of: "'", with: "''")
        let reader: String
        switch url.pathExtension.lowercased() {
        case "parquet", "pq":
            reader = "read_parquet('\(escaped)')"
        case "csv", "tsv", "txt":
            reader = "read_csv('\(escaped)')"
        case "json", "ndjson", "jsonl":
            reader = "read_json('\(escaped)')"
        case "arrow", "feather", "ipc":
            reader = "read_ipc('\(escaped)')"
        case let ext:
            throw ZQDataError("unsupported file format '.\(ext)'")
        }
        return try dataFrame(fromQuery: "SELECT * FROM \(reader)")
    }

    /// Runs an arbitrary SQL query and materializes the result.
    ///
    /// Two-phase: a `LIMIT 0` probe through duckdb-swift discovers names
    /// and types, then the full query — with every column cast to DOUBLE
    /// or VARCHAR in SQL — extracts through the C-API bulk reader.
    /// duckdb-swift's per-element access rescans every chunk per read
    /// (O(rows × chunks)); the bulk path is a straight memcpy per numeric
    /// column, ~200× faster at 10M rows.
    public func dataFrame(fromQuery sql: String) throws -> ZQDataFrame {
        let probe: ResultSet
        do {
            probe = try connection.query("SELECT * FROM (\(sql)) __q LIMIT 0")
        } catch {
            throw ZQDataError("query failed: \(error)")
        }

        var specs: [DuckDBFastReader.ColumnSpec] = []
        var selections: [String] = []
        for index in 0..<probe.columnCount {
            let name = probe.columnName(at: index)
            let quoted = "\"\(name.replacingOccurrences(of: "\"", with: "\"\""))\""
            switch probe.column(at: index).underlyingDatabaseType {
            case .double, .float, .decimal,
                 .tinyint, .smallint, .integer, .bigint,
                 .utinyint, .usmallint, .uinteger, .ubigint:
                specs.append(.init(name: name, isNumeric: true))
                selections.append("CAST(\(quoted) AS DOUBLE) AS \(quoted)")
            case .boolean:
                specs.append(.init(name: name, isNumeric: true))
                selections.append("CAST(CAST(\(quoted) AS INTEGER) AS DOUBLE) AS \(quoted)")
            default:
                // Text, dates, and everything else: string rendering.
                specs.append(.init(name: name, isNumeric: false))
                selections.append("CAST(\(quoted) AS VARCHAR) AS \(quoted)")
            }
        }
        guard !specs.isEmpty else { return try ZQDataFrame() }

        let castSQL = "SELECT \(selections.joined(separator: ", ")) FROM (\(sql)) __q"
        let reader = try DuckDBFastReader()
        return try reader.read(sql: castSQL, columns: specs)
    }

    // MARK: - Saving

    /// Writes a data frame to disk via DuckDB `COPY`. Format inferred from
    /// the extension (.parquet or .csv).
    public func save(_ frame: ZQDataFrame, to url: URL) throws {
        if url.pathExtension.lowercased() == "dta" {
            try DTAWriter.write(frame, to: url)
            return
        }

        let format: String
        switch url.pathExtension.lowercased() {
        case "parquet", "pq": format = "PARQUET"
        case "csv": format = "CSV, HEADER"
        case let ext: throw ZQDataError("unsupported save format '.\(ext)'")
        }

        try register(frame, as: "__metrika_save")
        defer { _ = try? connection.query("DROP TABLE IF EXISTS __metrika_save") }
        let escaped = url.path.replacingOccurrences(of: "'", with: "''")
        _ = try connection.query(
            "COPY __metrika_save TO '\(escaped)' (FORMAT \(format))"
        )
    }

    /// Registers a data frame as a DuckDB table (used by save and by SQL
    /// pushdown of `gen`/`collapse` expressions).
    private func register(_ frame: ZQDataFrame, as table: String) throws {
        _ = try? connection.query("DROP TABLE IF EXISTS \(table)")
        let declarations = frame.columns.map { column in
            let type = column.data.isNumeric ? "DOUBLE" : "VARCHAR"
            return "\"\(column.name)\" \(type)"
        }
        guard !declarations.isEmpty else {
            throw ZQDataError("nothing to save: dataset has no variables")
        }
        _ = try connection.query(
            "CREATE TABLE \(table) (\(declarations.joined(separator: ", ")))"
        )

        let appender = try Appender(connection: connection, table: table)
        for row in 0..<frame.rowCount {
            for column in frame.columns {
                switch column.data {
                case .float64(let values, let missing):
                    try appender.append(missing[row] ? nil : values[row])
                case .string(let values):
                    try appender.append(values[row])
                }
            }
            try appender.endRow()
        }
        try appender.flush()
    }
}