import SwiftUI

struct JSONTreeView: View {
    let data: Data

    @State private var root: JSONNode?

    var body: some View {
        Group {
            if let root {
                JSONNodeExpandableView(
                    node: root,
                    collapseByDefault: false,
                    childCollapseThreshold: 50
                )
            } else {
                Text("Invalid JSON")
                    .font(.system(.caption, design: .monospaced))
                    .foregroundStyle(.secondary)
            }
        }
        .task {
            root = parseJSON(data)
        }
    }

    private func parseJSON(_ data: Data) -> JSONNode? {
        guard let obj = try? JSONSerialization.jsonObject(with: data) else { return nil }
        return buildNode(key: nil, value: obj)
    }

    private func buildNode(key: String?, value: Any) -> JSONNode {
        if let dict = value as? [String: Any] {
            let children = dict.sorted(by: { $0.key < $1.key }).map { buildNode(key: $0.key, value: $0.value) }
            return JSONNode(key: key, kind: .object, displayValue: "{\(dict.count)}", children: children)
        } else if let arr = value as? [Any] {
            let children = arr.enumerated().map { buildNode(key: "[\($0.offset)]", value: $0.element) }
            return JSONNode(key: key, kind: .array, displayValue: "[\(arr.count)]", children: children)
        } else if let str = value as? String {
            return JSONNode(key: key, kind: .string, displayValue: "\"\(str)\"", children: [])
        } else if let num = value as? NSNumber {
            if CFBooleanGetTypeID() == CFGetTypeID(num) {
                let boolVal = num.boolValue
                return JSONNode(key: key, kind: .bool, displayValue: boolVal ? "true" : "false", children: [])
            }
            return JSONNode(key: key, kind: .number, displayValue: "\(num)", children: [])
        } else if value is NSNull {
            return JSONNode(key: key, kind: .null, displayValue: "null", children: [])
        } else {
            return JSONNode(key: key, kind: .string, displayValue: "\(value)", children: [])
        }
    }
}

// MARK: - Model

private struct JSONNode: Identifiable {
    let id = UUID()
    let key: String?
    let kind: JSONValueKind
    let displayValue: String
    let children: [JSONNode]

    var childCount: Int {
        children.count + children.reduce(0) { $0 + $1.childCount }
    }
}

private enum JSONValueKind {
    case object, array, string, number, bool, null
}

// MARK: - Expandable wrapper (allows per-node expand state)

private struct JSONNodeExpandableView: View {
    let node: JSONNode
    let collapseByDefault: Bool
    let childCollapseThreshold: Int

    @State private var isExpanded: Bool

    init(node: JSONNode, collapseByDefault: Bool, childCollapseThreshold: Int) {
        self.node = node
        self.collapseByDefault = collapseByDefault
        self.childCollapseThreshold = childCollapseThreshold
        _isExpanded = State(initialValue: !collapseByDefault)
    }

    var body: some View {
        if node.children.isEmpty {
            leafRow
        } else {
            DisclosureGroup(isExpanded: $isExpanded) {
                ForEach(node.children) { child in
                    JSONNodeExpandableView(
                        node: child,
                        collapseByDefault: child.childCount > childCollapseThreshold,
                        childCollapseThreshold: childCollapseThreshold
                    )
                }
            } label: {
                labelRow(valueColor: .secondary)
            }
        }
    }

    private var leafRow: some View {
        labelRow(valueColor: valueColor)
    }

    private func labelRow(valueColor: Color) -> some View {
        HStack(alignment: .firstTextBaseline, spacing: 4) {
            if let key = node.key {
                Text(key + ":")
                    .font(.system(.caption, design: .monospaced))
                    .bold()
                    .foregroundStyle(.primary)
            }

            Text(node.displayValue)
                .font(.system(.caption, design: .monospaced))
                .foregroundStyle(valueColor)
                .textSelection(.enabled)
        }
    }

    private var valueColor: Color {
        switch node.kind {
        case .string: return .green
        case .number: return .blue
        case .bool: return .orange
        case .null: return .gray
        case .object, .array: return .secondary
        }
    }
}