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

/// Column payload. Numeric data is Float64 end-to-end (CLAUDE.md §5: all
/// CPU compute in Float64) with an explicit validity mask mirroring Arrow
/// semantics: `missing[i] == true` means observation `i` is missing, and
/// the corresponding value slot is undefined.
public enum ZQColumnData: Equatable, Sendable {
    case float64(values: [Double], missing: [Bool])
    case string([String?])

    /// Value slots at missing positions are undefined (often NaN), so
    /// equality compares masks and only the valid values — synthesized
    /// equality would fail on NaN ≠ NaN.
    public static func == (lhs: ZQColumnData, rhs: ZQColumnData) -> Bool {
        switch (lhs, rhs) {
        case let (.float64(leftValues, leftMissing), .float64(rightValues, rightMissing)):
            guard leftMissing == rightMissing,
                  leftValues.count == rightValues.count else { return false }
            for i in 0..<leftValues.count where !leftMissing[i] {
                if leftValues[i] != rightValues[i] { return false }
            }
            return true
        case let (.string(left), .string(right)):
            return left == right
        default:
            return false
        }
    }

    public var count: Int {
        switch self {
        case .float64(let values, _): return values.count
        case .string(let values): return values.count
        }
    }

    public var missingCount: Int {
        switch self {
        case .float64(_, let missing): return missing.count { $0 }
        case .string(let values): return values.count { $0 == nil }
        }
    }

    public var isNumeric: Bool {
        if case .float64 = self { return true }
        return false
    }
}

/// A named column of a data frame.
public struct ZQColumn: Equatable, Sendable {
    public var name: String
    public var data: ZQColumnData

    public init(name: String, data: ZQColumnData) {
        self.name = name
        self.data = data
    }

    public var count: Int { data.count }

    /// Numeric view: values with their validity mask, or nil for string
    /// columns.
    public var numeric: (values: [Double], missing: [Bool])? {
        if case .float64(let values, let missing) = data {
            return (values, missing)
        }
        return nil
    }
}

public struct ZQDataError: Error, Equatable, Sendable, CustomStringConvertible {
    public var message: String
    public init(_ message: String) { self.message = message }
    public var description: String { message }
}

/// In-memory data frame: the working dataset of a session. Materialized
/// from DuckDB query results; columnar, immutable value semantics.
public struct ZQDataFrame: Equatable, Sendable {
    public private(set) var columns: [ZQColumn]
    public let rowCount: Int

    public init(columns: [ZQColumn] = []) throws(ZQDataError) {
        let counts = Set(columns.map(\.count))
        guard counts.count <= 1 else {
            throw ZQDataError("columns have inconsistent lengths: \(counts.sorted())")
        }
        self.columns = columns
        self.rowCount = counts.first ?? 0
    }

    public var isEmpty: Bool { columns.isEmpty || rowCount == 0 }
    public var columnNames: [String] { columns.map(\.name) }

    public func column(named name: String) -> ZQColumn? {
        columns.first { $0.name == name }
    }

    public func requireColumn(_ name: String) throws(ZQDataError) -> ZQColumn {
        guard let column = column(named: name) else {
            throw ZQDataError("variable '\(name)' not found")
        }
        return column
    }

    public func requireNumeric(
        _ name: String
    ) throws(ZQDataError) -> (values: [Double], missing: [Bool]) {
        guard let numeric = try requireColumn(name).numeric else {
            throw ZQDataError("variable '\(name)' is a string variable")
        }
        return numeric
    }

    /// Adds a new column; fails if the name already exists (Stata `generate`
    /// semantics — use `replaceColumn` for `replace`).
    public mutating func addColumn(_ column: ZQColumn) throws(ZQDataError) {
        guard self.column(named: column.name) == nil else {
            throw ZQDataError("variable '\(column.name)' already defined")
        }
        guard columns.isEmpty || column.count == rowCount else {
            throw ZQDataError(
                "column '\(column.name)' has \(column.count) rows, expected \(rowCount)"
            )
        }
        columns.append(column)
    }

    /// Replaces an existing column in place (Stata `replace`).
    public mutating func replaceColumn(_ column: ZQColumn) throws(ZQDataError) {
        guard let index = columns.firstIndex(where: { $0.name == column.name }) else {
            throw ZQDataError("variable '\(column.name)' not found")
        }
        guard column.count == rowCount else {
            throw ZQDataError(
                "column '\(column.name)' has \(column.count) rows, expected \(rowCount)"
            )
        }
        columns[index] = column
    }

    public mutating func dropColumns(_ names: [String]) throws(ZQDataError) {
        for name in names {
            guard columns.contains(where: { $0.name == name }) else {
                throw ZQDataError("variable '\(name)' not found")
            }
        }
        columns.removeAll { names.contains($0.name) }
    }

    /// Row subset by keep-mask (`mask[i] == true` keeps observation `i`).
    public func filtered(by mask: [Bool]) throws(ZQDataError) -> ZQDataFrame {
        guard mask.count == rowCount else {
            throw ZQDataError("filter mask length \(mask.count) != \(rowCount) rows")
        }
        let kept = (0..<rowCount).filter { mask[$0] }
        let newColumns = columns.map { column -> ZQColumn in
            switch column.data {
            case .float64(let values, let missing):
                return ZQColumn(
                    name: column.name,
                    data: .float64(
                        values: kept.map { values[$0] },
                        missing: kept.map { missing[$0] }
                    )
                )
            case .string(let values):
                return ZQColumn(name: column.name, data: .string(kept.map { values[$0] }))
            }
        }
        // Lengths are uniform by construction.
        return try! ZQDataFrame(columns: newColumns)
    }
}