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

import Foundation

/// Native writer for Stata .dta format 118 (Stata 14+, UTF-8, little
/// endian). Numeric columns store as double; string columns as str# up to
/// 2045 bytes. The map section carries real offsets, so Stata, haven, and
/// pandas can seek.
enum DTAWriter {

    static func write(_ frame: ZQDataFrame, to url: URL) throws {
        try data(for: frame).write(to: url)
    }

    static func data(for frame: ZQDataFrame) throws -> Data {
        guard !frame.columns.isEmpty else {
            throw ZQDataError("nothing to save: dataset has no variables")
        }
        let variableCount = frame.columns.count
        guard variableCount <= 32767 else {
            throw ZQDataError(".dta 118 supports at most 32,767 variables")
        }

        // str# width per string column (UTF-8 bytes, Stata counts bytes).
        var stringWidths = [Int](repeating: 0, count: variableCount)
        for (j, column) in frame.columns.enumerated() {
            if case .string(let values) = column.data {
                var width = 1
                for value in values {
                    if let value {
                        let bytes = value.utf8.count
                        guard bytes <= DTAFormat.maxFixedStringLength else {
                            throw ZQDataError(
                                "string variable '\(column.name)' exceeds 2045 bytes — strL writing not supported yet"
                            )
                        }
                        width = max(width, bytes)
                    }
                }
                stringWidths[j] = width
            }
        }

        var out = Data()
        var mapOffsets = [UInt64](repeating: 0, count: 14)

        func append(_ text: String) { out.append(Data(text.utf8)) }
        func appendUInt(_ value: UInt64, _ byteCount: Int) {
            for shift in 0..<byteCount {
                out.append(UInt8(truncatingIfNeeded: value >> (8 * shift)))
            }
        }
        /// Null-terminated field (names, formats): content ≤ width−1 bytes.
        func appendPadded(_ text: String, _ width: Int) {
            var bytes = Data(text.utf8).prefix(width - 1)
            bytes.append(contentsOf: [UInt8](repeating: 0, count: width - bytes.count))
            out.append(bytes)
        }

        /// str# data cell: exactly width bytes, null-padded, no terminator
        /// required when the content fills the field.
        func appendCell(_ text: String, _ width: Int) {
            var bytes = Data(text.utf8).prefix(width)
            bytes.append(contentsOf: [UInt8](repeating: 0, count: width - bytes.count))
            out.append(bytes)
        }

        // Header.
        mapOffsets[0] = 0
        append("<stata_dta><header><release>118</release>")
        append("<byteorder>LSF</byteorder><K>")
        appendUInt(UInt64(variableCount), 2)
        append("</K><N>")
        appendUInt(UInt64(frame.rowCount), 8)
        append("</N><label>")
        appendUInt(0, 2)
        append("</label><timestamp>")
        appendUInt(0, 1)
        append("</timestamp></header>")

        // Map: placeholder now, patched at the end.
        mapOffsets[1] = UInt64(out.count)
        append("<map>")
        let mapPayloadOffset = out.count
        out.append(Data(repeating: 0, count: 14 * 8))
        append("</map>")

        // Variable types.
        mapOffsets[2] = UInt64(out.count)
        append("<variable_types>")
        for (j, column) in frame.columns.enumerated() {
            let code: UInt16 = column.data.isNumeric
                ? DTAFormat.typeDouble
                : UInt16(stringWidths[j])
            appendUInt(UInt64(code), 2)
        }
        append("</variable_types>")

        // Names (129-byte fields).
        mapOffsets[3] = UInt64(out.count)
        append("<varnames>")
        for column in frame.columns {
            guard column.name.utf8.count <= 128 else {
                throw ZQDataError("variable name '\(column.name)' exceeds 128 bytes")
            }
            appendPadded(column.name, 129)
        }
        append("</varnames>")

        // Sort list: unsorted.
        mapOffsets[4] = UInt64(out.count)
        append("<sortlist>")
        out.append(Data(repeating: 0, count: 2 * (variableCount + 1)))
        append("</sortlist>")

        // Display formats (57-byte fields).
        mapOffsets[5] = UInt64(out.count)
        append("<formats>")
        for (j, column) in frame.columns.enumerated() {
            let format = column.data.isNumeric ? "%10.0g" : "%\(stringWidths[j])s"
            appendPadded(format, 57)
        }
        append("</formats>")

        // No value labels or variable labels.
        mapOffsets[6] = UInt64(out.count)
        append("<value_label_names>")
        out.append(Data(repeating: 0, count: 129 * variableCount))
        append("</value_label_names>")

        mapOffsets[7] = UInt64(out.count)
        append("<variable_labels>")
        out.append(Data(repeating: 0, count: 321 * variableCount))
        append("</variable_labels>")

        mapOffsets[8] = UInt64(out.count)
        append("<characteristics></characteristics>")

        // Data, row-major.
        mapOffsets[9] = UInt64(out.count)
        append("<data>")
        for row in 0..<frame.rowCount {
            for (j, column) in frame.columns.enumerated() {
                switch column.data {
                case .float64(let values, let missing):
                    let value = missing[row] ? DTAFormat.doubleMissingValue : values[row]
                    appendUInt(value.bitPattern, 8)
                case .string(let values):
                    // Stata's string missing is the empty string.
                    appendCell(values[row] ?? "", stringWidths[j])
                }
            }
        }
        append("</data>")

        mapOffsets[10] = UInt64(out.count)
        append("<strls></strls>")

        mapOffsets[11] = UInt64(out.count)
        append("<value_labels></value_labels>")

        mapOffsets[12] = UInt64(out.count)
        append("</stata_dta>")
        mapOffsets[13] = UInt64(out.count)

        // Patch the map.
        for (index, offset) in mapOffsets.enumerated() {
            for shift in 0..<8 {
                out[mapPayloadOffset + index * 8 + shift] =
                    UInt8(truncatingIfNeeded: offset >> (8 * shift))
            }
        }
        return out
    }
}