ReadViewSDK/Sources/RDReaderView/EPUBCore/RDEPUBParser+TOC.swift

import Foundation

extension RDEPUBParser {

    func parseTOC(from packageDocument: OPFPackageDocument, opfURL: URL) -> [EPUBTableOfContentsItem] {
        let opfDirectoryURL = opfURL.deletingLastPathComponent()

        if let ncxItem = packageDocument.ncxItem ?? packageDocument.manifest.first(where: { $0.isNCX }) {
            let ncxURL = opfDirectoryURL.appendingPathComponent(ncxItem.href)
            let items = parseNCXDocumentItems(at: ncxURL)
            if !items.isEmpty {
                return items
            }
        }

        if let navItem = packageDocument.navigationItem {
            let navURL = opfDirectoryURL.appendingPathComponent(navItem.href)
            let items = parseNavDocumentItems(at: navURL, baseURL: navURL.deletingLastPathComponent())
            if !items.isEmpty {
                return items
            }
        }

        return spine.map { EPUBTableOfContentsItem(title: $0.title, href: $0.href) }
    }

    func parseNCXDocumentItems(at ncxURL: URL) -> [EPUBTableOfContentsItem] {
        guard FileManager.default.fileExists(atPath: ncxURL.path),
              let parser = XMLParser(contentsOf: ncxURL) else {
            return []
        }

        let delegate = NCXParserDelegate(baseURL: ncxURL.deletingLastPathComponent()) { [weak self] href, baseURL in
            self?.normalizeTOCHref(href, relativeTo: baseURL)
        }
        parser.shouldProcessNamespaces = false
        parser.delegate = delegate
        guard parser.parse() else {
            return []
        }
        return delegate.items
    }

    func parseNavDocumentItems(at navURL: URL, baseURL: URL) -> [EPUBTableOfContentsItem] {
        guard FileManager.default.fileExists(atPath: navURL.path),
              let parser = XMLParser(contentsOf: navURL) else {
            return []
        }

        let delegate = NavDocumentParserDelegate(baseURL: baseURL) { [weak self] href, currentBaseURL in
            self?.normalizeTOCHref(href, relativeTo: currentBaseURL)
        }
        parser.shouldProcessNamespaces = false
        parser.delegate = delegate
        guard parser.parse() else {
            return []
        }
        return delegate.items
    }

    func normalizeTOCHref(_ href: String, relativeTo baseURL: URL) -> String? {
        let components = href.split(separator: "#", maxSplits: 1, omittingEmptySubsequences: false)
        let pathPart = components.first.map(String.init) ?? href
        let fragment = components.count > 1 ? String(components[1]) : nil

        let normalizedPath: String
        if pathPart.isEmpty {
            normalizedPath = ""
        } else {
            normalizedPath = normalize(href: pathPart, relativeTo: baseURL)
        }

        if let fragment, !fragment.isEmpty {
            if normalizedPath.isEmpty {
                return "#\(fragment)"
            }
            return "\(normalizedPath)#\(fragment)"
        }

        return normalizedPath.isEmpty ? nil : normalizedPath
    }
}

private final class TOCNode {
    var title = ""
    var href: String?
    var children: [TOCNode] = []

    func asItem() -> EPUBTableOfContentsItem? {
        let trimmedTitle = title.trimmingCharacters(in: .whitespacesAndNewlines)
        guard let href, !href.isEmpty else {
            return nil
        }
        return EPUBTableOfContentsItem(
            title: trimmedTitle.isEmpty ? href : trimmedTitle,
            href: href,
            children: children.compactMap { $0.asItem() }
        )
    }
}

private final class TOCTreeBuilder {

    private var nodeStack: [TOCNode] = []

    private var rootNodes: [TOCNode] = []

    var items: [EPUBTableOfContentsItem] {
        rootNodes.compactMap { $0.asItem() }
    }

    func beginNode() {
        nodeStack.append(TOCNode())
    }

    func updateCurrentHref(_ href: String?) {
        nodeStack.last?.href = href
    }

    func appendCurrentTitle(_ text: String) {
        nodeStack.last?.title += text
    }

    func endNode() {
        guard let node = nodeStack.popLast() else {
            return
        }
        if let parent = nodeStack.last {
            parent.children.append(node)
        } else {
            rootNodes.append(node)
        }
    }

    func reset() {
        nodeStack.removeAll()
    }
}

private final class NCXParserDelegate: NSObject, XMLParserDelegate {
    private let baseURL: URL
    private let hrefNormalizer: (String, URL) -> String?
    private let treeBuilder = TOCTreeBuilder()
    private var readingLabelText = false

    init(baseURL: URL, hrefNormalizer: @escaping (String, URL) -> String?) {
        self.baseURL = baseURL
        self.hrefNormalizer = hrefNormalizer
    }

    var items: [EPUBTableOfContentsItem] {
        treeBuilder.items
    }

    func parser(
        _ parser: XMLParser,
        didStartElement elementName: String,
        namespaceURI: String?,
        qualifiedName qName: String?,
        attributes attributeDict: [String: String] = [:]
    ) {
        let name = XMLName.localName(from: qName ?? elementName)

        switch name {
        case "navPoint":
            treeBuilder.beginNode()
        case "content":
            guard let src = attributeDict["src"] else {
                return
            }
            treeBuilder.updateCurrentHref(hrefNormalizer(src, baseURL))
        case "text":
            readingLabelText = true
        default:
            break
        }
    }

    func parser(_ parser: XMLParser, foundCharacters string: String) {
        if readingLabelText {
            treeBuilder.appendCurrentTitle(string)
        }
    }

    func parser(
        _ parser: XMLParser,
        didEndElement elementName: String,
        namespaceURI: String?,
        qualifiedName qName: String?
    ) {
        let name = XMLName.localName(from: qName ?? elementName)

        switch name {
        case "text":
            readingLabelText = false
        case "navPoint":
            treeBuilder.endNode()
        default:
            break
        }
    }
}

private final class NavDocumentParserDelegate: NSObject, XMLParserDelegate {
    private let baseURL: URL
    private let hrefNormalizer: (String, URL) -> String?
    private var insideTOCNav = false
    private var listDepth = 0
    private let treeBuilder = TOCTreeBuilder()
    private var collectingAnchorText = false

    init(baseURL: URL, hrefNormalizer: @escaping (String, URL) -> String?) {
        self.baseURL = baseURL
        self.hrefNormalizer = hrefNormalizer
    }

    var items: [EPUBTableOfContentsItem] {
        treeBuilder.items
    }

    func parser(
        _ parser: XMLParser,
        didStartElement elementName: String,
        namespaceURI: String?,
        qualifiedName qName: String?,
        attributes attributeDict: [String: String] = [:]
    ) {
        let name = XMLName.localName(from: qName ?? elementName)

        switch name {
        case "nav":
            insideTOCNav = isTOCNav(attributes: attributeDict)
        case "ol":
            if insideTOCNav {
                listDepth += 1
            }
        case "li":
            if insideTOCNav, listDepth > 0 {
                treeBuilder.beginNode()
            }
        case "a":
            guard insideTOCNav else {
                return
            }
            collectingAnchorText = true
            if let href = attributeDict["href"] {
                treeBuilder.updateCurrentHref(hrefNormalizer(href, baseURL))
            }
        default:
            break
        }
    }

    func parser(_ parser: XMLParser, foundCharacters string: String) {
        if collectingAnchorText {
            treeBuilder.appendCurrentTitle(string)
        }
    }

    func parser(
        _ parser: XMLParser,
        didEndElement elementName: String,
        namespaceURI: String?,
        qualifiedName qName: String?
    ) {
        let name = XMLName.localName(from: qName ?? elementName)

        switch name {
        case "a":
            collectingAnchorText = false
        case "li":
            guard insideTOCNav else {
                break
            }
            treeBuilder.endNode()
        case "ol":
            if insideTOCNav {
                listDepth = max(0, listDepth - 1)
            }
        case "nav":
            insideTOCNav = false
            listDepth = 0
            treeBuilder.reset()
            collectingAnchorText = false
        default:
            break
        }
    }

    private func isTOCNav(attributes: [String: String]) -> Bool {
        if let epubType = attributes["epub:type"]?.lowercased(), epubType.contains("toc") {
            return true
        }
        if let type = attributes["type"]?.lowercased(), type.contains("toc") {
            return true
        }
        if let role = attributes["role"]?.lowercased(), role.contains("doc-toc") {
            return true
        }
        return false
    }
}