import Foundation

public protocol MarkdownRendering: Sendable {
    func blocks(for markdown: String) -> [RenderedMarkdownBlock]
    func inlineMarkdown(for source: String) -> AttributedString
}

public enum RenderedMarkdownBlock: Identifiable, Hashable, Sendable {
    case heading(id: UUID, level: Int, text: String)
    case paragraph(id: UUID, text: AttributedString)
    case codeBlock(id: UUID, language: String?, code: String)
    case task(id: UUID, checked: Bool, text: AttributedString)
    case image(id: UUID, altText: String, source: String)
    case blank(id: UUID)

    public var id: UUID {
        switch self {
        case .heading(let id, _, _): id
        case .paragraph(let id, _): id
        case .codeBlock(let id, _, _): id
        case .task(let id, _, _): id
        case .image(let id, _, _): id
        case .blank(let id): id
        }
    }
}

public struct MarkdownRenderer: MarkdownRendering {
    public init() {}

    public func blocks(for markdown: String) -> [RenderedMarkdownBlock] {
        var blocks: [RenderedMarkdownBlock] = []
        var codeLines: [String] = []
        var codeLanguage: String?
        var isInCodeBlock = false

        for line in markdown.split(separator: "\n", omittingEmptySubsequences: false).map(String.init) {
            if line.hasPrefix("```") {
                if isInCodeBlock {
                    blocks.append(.codeBlock(id: UUID(), language: codeLanguage, code: codeLines.joined(separator: "\n")))
                    codeLines.removeAll()
                    codeLanguage = nil
                    isInCodeBlock = false
                } else {
                    isInCodeBlock = true
                    codeLanguage = String(line.dropFirst(3)).nilIfEmpty
                }
                continue
            }

            if isInCodeBlock {
                codeLines.append(line)
                continue
            }

            blocks.append(block(for: line))
        }

        if isInCodeBlock {
            blocks.append(.codeBlock(id: UUID(), language: codeLanguage, code: codeLines.joined(separator: "\n")))
        }

        return blocks
    }

    public func inlineMarkdown(for source: String) -> AttributedString {
        (try? AttributedString(markdown: source)) ?? AttributedString(source)
    }

    private func block(for line: String) -> RenderedMarkdownBlock {
        let trimmed = line.trimmingCharacters(in: .whitespaces)

        guard !trimmed.isEmpty else {
            return .blank(id: UUID())
        }

        if let heading = parseHeading(trimmed) {
            return .heading(id: UUID(), level: heading.level, text: heading.text)
        }

        if let task = parseTask(trimmed) {
            return .task(id: UUID(), checked: task.checked, text: inlineMarkdown(for: task.text))
        }

        if let image = parseImage(trimmed) {
            return .image(id: UUID(), altText: image.altText, source: image.source)
        }

        return .paragraph(id: UUID(), text: inlineMarkdown(for: line))
    }

    private func parseHeading(_ line: String) -> (level: Int, text: String)? {
        let markerCount = line.prefix { $0 == "#" }.count
        guard (1...6).contains(markerCount), line.dropFirst(markerCount).first == " " else {
            return nil
        }
        return (markerCount, String(line.dropFirst(markerCount + 1)))
    }

    private func parseTask(_ line: String) -> (checked: Bool, text: String)? {
        if line.hasPrefix("- [x] ") || line.hasPrefix("- [X] ") {
            return (true, String(line.dropFirst(6)))
        }
        if line.hasPrefix("- [ ] ") {
            return (false, String(line.dropFirst(6)))
        }
        return nil
    }

    private func parseImage(_ line: String) -> (altText: String, source: String)? {
        guard line.hasPrefix("!["), let closeAlt = line.firstIndex(of: "]") else { return nil }
        let afterAlt = line[line.index(after: closeAlt)...]
        guard afterAlt.hasPrefix("("), afterAlt.hasSuffix(")") else { return nil }
        let alt = String(line[line.index(line.startIndex, offsetBy: 2)..<closeAlt])
        let source = String(afterAlt.dropFirst().dropLast())
        return (alt, source)
    }
}

private extension String {
    var nilIfEmpty: String? {
        isEmpty ? nil : self
    }
}