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

import Foundation

/// Native reader for Stata .dta formats 117 (Stata 13), 118 (Stata 14–18)
/// and 119 (>32k variables). Numeric storage types widen to Float64 with
/// Stata missing codes (`.`, `.a`–`.z`) mapped to the validity mask; str#
/// and strL variables load as strings. Value labels are not yet applied.
enum DTAReader {

    static func read(contentsOf url: URL) throws -> ZQDataFrame {
        let data = try Data(contentsOf: url)
        return try read(data: data)
    }

    static func read(data: Data) throws -> ZQDataFrame {
        var cursor = DTACursor(data: data)

        guard data.starts(with: Data("<stata_dta>".utf8)) else {
            // Pre-117 files begin with a raw version byte (0x66–0x73).
            if let first = data.first, (0x66...0x76).contains(first) {
                throw ZQDataError(
                    ".dta formats older than 117 (Stata 13) are not supported — re-save with a modern Stata"
                )
            }
            throw ZQDataError("not a Stata .dta file")
        }

        // Header.
        try cursor.expect("<stata_dta><header><release>")
        let releaseText = try cursor.readPaddedString(3, release: 117)
        guard let release = Int(releaseText) else {
            throw ZQDataError("corrupt .dta: bad release '\(releaseText)'")
        }
        guard let layout = DTAFormat.Layout(release: release) else {
            throw ZQDataError(
                ".dta format \(release) is not supported (117–119 are)"
            )
        }
        try cursor.expect("</release><byteorder>")
        let byteOrder = try cursor.readPaddedString(3, release: 117)
        cursor.bigEndian = byteOrder == "MSF"
        try cursor.expect("</byteorder><K>")
        let variableCount = Int(try cursor.readUInt(layout.variableCountBytes))
        try cursor.expect("</K><N>")
        let observationCount = Int(try cursor.readUInt(layout.observationCountBytes))
        try cursor.expect("</N><label>")
        let labelLength = Int(try cursor.readUInt(layout.datasetLabelLengthBytes))
        _ = try cursor.readBytes(labelLength)
        try cursor.expect("</label><timestamp>")
        let timestampLength = Int(try cursor.readUInt(1))
        _ = try cursor.readBytes(timestampLength)
        try cursor.expect("</timestamp></header>")

        // Map: 14 file offsets. Trusted for jumping to <strls> later; the
        // sections up to <data> are read sequentially.
        try cursor.expect("<map>")
        var map: [UInt64] = []
        for _ in 0..<14 { map.append(try cursor.readUInt(8)) }
        try cursor.expect("</map>")

        // Variable descriptors.
        try cursor.expect("<variable_types>")
        var types: [UInt16] = []
        for _ in 0..<variableCount {
            types.append(UInt16(truncatingIfNeeded: try cursor.readUInt(2)))
        }
        try cursor.expect("</variable_types>")

        try cursor.expect("<varnames>")
        var names: [String] = []
        for _ in 0..<variableCount {
            names.append(try cursor.readPaddedString(layout.nameBytes, release: release))
        }
        try cursor.expect("</varnames>")

        try cursor.expect("<sortlist>")
        _ = try cursor.readBytes(layout.sortEntryBytes * (variableCount + 1))
        try cursor.expect("</sortlist>")

        try cursor.expect("<formats>")
        _ = try cursor.readBytes(layout.formatBytes * variableCount)
        try cursor.expect("</formats>")

        try cursor.expect("<value_label_names>")
        _ = try cursor.readBytes(layout.labelNameBytes * variableCount)
        try cursor.expect("</value_label_names>")

        try cursor.expect("<variable_labels>")
        _ = try cursor.readBytes(layout.variableLabelBytes * variableCount)
        try cursor.expect("</variable_labels>")

        // Characteristics: <ch> length-prefixed blobs until the closer.
        try cursor.expect("<characteristics>")
        while cursor.remaining >= 4 {
            let peek = try DTACursor(data: data).peekBytes(at: cursor.offset, count: 4)
            if peek == Data("</ch".utf8) { break }
            try cursor.expect("<ch>")
            let length = Int(try cursor.readUInt(4))
            _ = try cursor.readBytes(length)
            try cursor.expect("</ch>")
        }
        try cursor.expect("</characteristics>")

        // strLs are stored after <data>; when any variable is strL, decode
        // that section first (via the map) so cells can resolve.
        var strLTable: [UInt64: String] = [:]
        if types.contains(DTAFormat.typeStrL) {
            strLTable = try readStrLs(
                data: data, at: Int(map[10]), layout: layout,
                bigEndian: cursor.bigEndian, release: release
            )
        }

        // Data.
        try cursor.expect("<data>")
        var numericValues = [[Double]](
            repeating: [Double](repeating: 0, count: observationCount),
            count: variableCount
        )
        var numericMissing = [[Bool]](
            repeating: [Bool](repeating: false, count: observationCount),
            count: variableCount
        )
        var stringValues = [[String?]](repeating: [], count: variableCount)
        for j in 0..<variableCount where isStringType(types[j]) {
            stringValues[j] = [String?](repeating: nil, count: observationCount)
        }

        for row in 0..<observationCount {
            for j in 0..<variableCount {
                let type = types[j]
                switch type {
                case DTAFormat.typeDouble:
                    let value = try cursor.readDouble()
                    if value.isNaN || value > DTAFormat.doubleMissingThreshold {
                        numericMissing[j][row] = true
                        numericValues[j][row] = .nan
                    } else {
                        numericValues[j][row] = value
                    }
                case DTAFormat.typeFloat:
                    let value = try cursor.readFloat()
                    if value.isNaN || value > DTAFormat.floatMissingThreshold {
                        numericMissing[j][row] = true
                        numericValues[j][row] = .nan
                    } else {
                        numericValues[j][row] = Double(value)
                    }
                case DTAFormat.typeLong:
                    let value = Int32(truncatingIfNeeded: try cursor.readUInt(4))
                    if value >= DTAFormat.longMissingThreshold {
                        numericMissing[j][row] = true
                        numericValues[j][row] = .nan
                    } else {
                        numericValues[j][row] = Double(value)
                    }
                case DTAFormat.typeInt:
                    let value = Int16(truncatingIfNeeded: try cursor.readUInt(2))
                    if value >= DTAFormat.intMissingThreshold {
                        numericMissing[j][row] = true
                        numericValues[j][row] = .nan
                    } else {
                        numericValues[j][row] = Double(value)
                    }
                case DTAFormat.typeByte:
                    let value = Int8(truncatingIfNeeded: try cursor.readUInt(1))
                    if value >= DTAFormat.byteMissingThreshold {
                        numericMissing[j][row] = true
                        numericValues[j][row] = .nan
                    } else {
                        numericValues[j][row] = Double(value)
                    }
                case DTAFormat.typeStrL:
                    let v = try cursor.readUInt(layout.strLVarBytes)
                    let o = try cursor.readUInt(layout.strLObsBytes)
                    if v == 0 && o == 0 {
                        stringValues[j][row] = nil
                    } else {
                        stringValues[j][row] = strLTable[strLKey(v: v, o: o)] ?? ""
                    }
                case 1...UInt16(DTAFormat.maxFixedStringLength):
                    let text = try cursor.readPaddedString(Int(type), release: release)
                    stringValues[j][row] = text.isEmpty ? nil : text
                default:
                    throw ZQDataError("corrupt .dta: unknown variable type \(type)")
                }
            }
        }
        try cursor.expect("</data>")

        // Assemble columns.
        var columns: [ZQColumn] = []
        for j in 0..<variableCount {
            if isStringType(types[j]) {
                columns.append(ZQColumn(name: names[j], data: .string(stringValues[j])))
            } else {
                columns.append(ZQColumn(
                    name: names[j],
                    data: .float64(values: numericValues[j], missing: numericMissing[j])
                ))
            }
        }
        return try ZQDataFrame(columns: columns)
    }

    // MARK: - Helpers

    private static func isStringType(_ type: UInt16) -> Bool {
        type == DTAFormat.typeStrL
            || (1...UInt16(DTAFormat.maxFixedStringLength)).contains(type)
    }

    private static func strLKey(v: UInt64, o: UInt64) -> UInt64 {
        // (v, o) packed; v < 2^16 in 118 and < 2^32 in 117/119, o < 2^48.
        (v << 48) | o
    }

    /// Parses the <strls> section: consecutive GSO records.
    private static func readStrLs(
        data: Data, at offset: Int, layout: DTAFormat.Layout,
        bigEndian: Bool, release: Int
    ) throws -> [UInt64: String] {
        var cursor = DTACursor(data: data)
        cursor.offset = offset
        cursor.bigEndian = bigEndian
        try cursor.expect("<strls>")

        var table: [UInt64: String] = [:]
        while cursor.remaining >= 3 {
            let marker = try cursor.peekBytes(at: cursor.offset, count: 3)
            if marker != Data("GSO".utf8) { break }
            _ = try cursor.readBytes(3)
            let v = try cursor.readUInt(4)
            let o = try cursor.readUInt(layout.gsoObsBytes)
            let kind = try cursor.readUInt(1)   // 129 binary, 130 ASCII/UTF-8
            let length = Int(try cursor.readUInt(4))
            var payload = try cursor.readBytes(length)
            if kind == 130, payload.last == 0 {
                payload = payload.dropLast()
            }
            let text = String(data: payload, encoding: .utf8)
                ?? String(data: payload, encoding: .isoLatin1)
                ?? ""
            table[strLKey(v: v, o: o)] = text
        }
        try cursor.expect("</strls>")
        return table
    }
}

extension DTACursor {
    /// Non-consuming read at an absolute offset.
    func peekBytes(at position: Int, count: Int) throws(ZQDataError) -> Data {
        guard position >= 0, position + count <= data.count else {
            throw ZQDataError("corrupt .dta: unexpected end of file at offset \(position)")
        }
        return data.subdata(in: position..<(position + count))
    }
}