import Foundation
import SwiftUI

/// A lightweight Markdown renderer for chat messages: fenced code blocks, headings,
/// bullet/numbered lists, and inline emphasis/links/`code`. Not a full CommonMark
/// implementation — just the constructs agents actually emit.
struct MarkdownText: View {
    let markdown: String
    /// Base prose size; every other size (headings, code, tables) is derived from
    /// it so a response renders on one consistent scale. Defaults to the transcript
    /// prose size (`SessionDetailView.transcriptFontSize`).
    var bodySize: CGFloat = 15

    var body: some View {
        VStack(alignment: .leading, spacing: 8) {
            ForEach(Array(Self.parse(markdown).enumerated()), id: \.offset) { _, block in
                switch block {
                case .code(let code):
                    codeBlock(code)
                case .text(let text):
                    textBlock(text)
                case .table(let rows):
                    tableView(rows)
                }
            }
        }
        // Anchor the whole response to one base size; inline code and body text
        // inherit it, so nothing drifts smaller than the prose around it.
        .font(.system(size: bodySize))
        // Soft off-white (vs. pure white) so chat prose reads calmly at night.
        .foregroundStyle(AppTheme.primaryText)
    }

    // MARK: - blocks

    private enum Block { case code(String), text(String), table([[String]]) }

    /// Parsed-block cache. The same message string is re-parsed on every body
    /// re-evaluation (scrolling, hover, a sibling row streaming) and on every chat
    /// reopen, yet structural parsing is independent of `bodySize` — so the source
    /// string is a complete key. Bounded; `NSCache` also evicts under memory pressure.
    private final class ParsedBlocks { let blocks: [Block]; init(_ b: [Block]) { self.blocks = b } }
    private static let blockCache: NSCache<NSString, ParsedBlocks> = {
        let cache = NSCache<NSString, ParsedBlocks>()
        cache.countLimit = 2048
        return cache
    }()

    private static func parse(_ markdown: String) -> [Block] {
        let key = markdown as NSString
        if let hit = blockCache.object(forKey: key) { return hit.blocks }
        let blocks = parseUncached(markdown)
        blockCache.setObject(ParsedBlocks(blocks), forKey: key)
        return blocks
    }

    private static func parseUncached(_ markdown: String) -> [Block] {
        var blocks: [Block] = []
        var textBuffer: [String] = []
        func flush() {
            let joined = textBuffer.joined(separator: "\n").trimmingCharacters(in: .whitespacesAndNewlines)
            if !joined.isEmpty { blocks.append(.text(joined)) }
            textBuffer = []
        }
        let lines = markdown.components(separatedBy: "\n")
        var index = 0
        while index < lines.count {
            if lines[index].trimmingCharacters(in: .whitespaces).hasPrefix("```") {
                flush()
                var code: [String] = []
                index += 1
                while index < lines.count,
                    !lines[index].trimmingCharacters(in: .whitespaces).hasPrefix("```") {
                    code.append(lines[index])
                    index += 1
                }
                blocks.append(.code(code.joined(separator: "\n")))
                index += 1  // consume closing fence
            } else if isTableStart(lines, index) {
                flush()
                var rows: [[String]] = [tableCells(lines[index])]
                index += 2  // header row + the |---|--- separator
                while index < lines.count, lines[index].contains("|"),
                    !lines[index].trimmingCharacters(in: .whitespaces).isEmpty {
                    rows.append(tableCells(lines[index]))
                    index += 1
                }
                blocks.append(.table(rows))
            } else {
                textBuffer.append(lines[index])
                index += 1
            }
        }
        flush()
        return blocks
    }

    // MARK: - tables

    /// A GitHub-style table: a `|`-bearing header line immediately followed by a
    /// `|---|:--:|` separator line.
    private static func isTableStart(_ lines: [String], _ index: Int) -> Bool {
        guard lines[index].contains("|"), index + 1 < lines.count else { return false }
        return isSeparatorRow(lines[index + 1])
    }

    private static func isSeparatorRow(_ line: String) -> Bool {
        let cells = tableCells(line)
        guard !cells.isEmpty else { return false }
        return cells.allSatisfy { cell in
            !cell.isEmpty && cell.allSatisfy { $0 == "-" || $0 == ":" } && cell.contains("-")
        }
    }

    /// Split a table row into trimmed cells, dropping the empties created by the
    /// leading/trailing pipes.
    private static func tableCells(_ line: String) -> [String] {
        var trimmed = line.trimmingCharacters(in: .whitespaces)
        if trimmed.hasPrefix("|") { trimmed.removeFirst() }
        if trimmed.hasSuffix("|") { trimmed.removeLast() }
        return trimmed.components(separatedBy: "|").map { $0.trimmingCharacters(in: .whitespaces) }
    }

    private func tableView(_ rows: [[String]]) -> some View {
        let columns = rows.map(\.count).max() ?? 0
        return Grid(alignment: .topLeading, horizontalSpacing: 14, verticalSpacing: 6) {
            ForEach(Array(rows.enumerated()), id: \.offset) { rowIndex, row in
                GridRow {
                    ForEach(0..<columns, id: \.self) { col in
                        inline(col < row.count ? row[col] : "")
                            .font(.system(size: bodySize, weight: rowIndex == 0 ? .semibold : .regular))
                            .frame(maxWidth: .infinity, alignment: .leading)
                    }
                }
                if rowIndex == 0, columns > 0 {
                    Divider().gridCellColumns(columns)
                }
            }
        }
        .padding(10)
        .background(.quaternary.opacity(0.25), in: .rect(cornerRadius: 8))
        .overlay(RoundedRectangle(cornerRadius: 8).strokeBorder(.quaternary, lineWidth: 1))
    }

    private func codeBlock(_ code: String) -> some View {
        ScrollView(.horizontal, showsIndicators: false) {
            Text(code)
                // Same point size as body/inline code, just monospaced — so code
                // doesn't shrink relative to the prose.
                .font(.system(size: bodySize, design: .monospaced))
                .textSelection(.enabled)
                .padding(8)
                .frame(maxWidth: .infinity, alignment: .leading)
        }
        .background(.quaternary.opacity(0.4), in: .rect(cornerRadius: 6))
    }

    private func textBlock(_ text: String) -> some View {
        VStack(alignment: .leading, spacing: 5) {
            ForEach(Array(text.components(separatedBy: "\n").enumerated()), id: \.offset) { _, raw in
                lineView(raw)
            }
        }
        // Breathing room between the visual lines of a single wrapped paragraph,
        // which the per-line VStack spacing above can't reach.
        .lineSpacing(4)
    }

    @ViewBuilder
    private func lineView(_ raw: String) -> some View {
        let trimmed = raw.trimmingCharacters(in: .whitespaces)
        if trimmed.isEmpty {
            Color.clear.frame(height: 3)
        } else if trimmed.hasPrefix("### ") {
            // Headings scale relative to the base prose size so the hierarchy holds
            // at any base and never collapses to the body size.
            inline(String(trimmed.dropFirst(4))).font(.system(size: bodySize * 1.13, weight: .semibold))
        } else if trimmed.hasPrefix("## ") {
            inline(String(trimmed.dropFirst(3))).font(.system(size: bodySize * 1.28, weight: .bold))
        } else if trimmed.hasPrefix("# ") {
            inline(String(trimmed.dropFirst(2))).font(.system(size: bodySize * 1.5, weight: .bold))
        } else if let bullet = Self.bulletContent(trimmed) {
            HStack(alignment: .firstTextBaseline, spacing: 6) {
                Text("•").foregroundStyle(.secondary)
                inline(bullet)
            }
        } else {
            inline(raw)
        }
    }

    /// Returns the content after a `- `, `* `, `+ ` or `N. ` list marker, else nil.
    private static func bulletContent(_ trimmed: String) -> String? {
        for marker in ["- ", "* ", "+ "] where trimmed.hasPrefix(marker) {
            return String(trimmed.dropFirst(marker.count))
        }
        // Numbered list: "12. text"
        let parts = trimmed.split(separator: " ", maxSplits: 1)
        if let first = parts.first, parts.count == 2,
            first.hasSuffix("."), Int(first.dropLast()) != nil {
            return String(parts[1])
        }
        return nil
    }

    private func inline(_ string: String) -> Text {
        Text(Self.attributedInline(string))
    }

    /// Inline-Markdown cache. `AttributedString(markdown:)` is the dominant per-row cost
    /// when a transcript first lays out — it runs once per prose line — and the same lines
    /// recur across re-renders and reopens, so memoize the parsed result. Independent of
    /// `bodySize` (callers apply the font), so the source string is a complete key.
    private final class InlineBox { let value: AttributedString; init(_ v: AttributedString) { self.value = v } }
    private static let inlineCache: NSCache<NSString, InlineBox> = {
        let cache = NSCache<NSString, InlineBox>()
        cache.countLimit = 16384
        return cache
    }()

    private static func attributedInline(_ string: String) -> AttributedString {
        let key = string as NSString
        if let hit = inlineCache.object(forKey: key) { return hit.value }
        let options = AttributedString.MarkdownParsingOptions(
            interpretedSyntax: .inlineOnlyPreservingWhitespace)
        // Fall back to the plain string on parse failure — `Text(AttributedString(string))`
        // renders identically to `Text(string)`, so callers see no behavioral change.
        let parsed = (try? AttributedString(markdown: string, options: options)) ?? AttributedString(string)
        inlineCache.setObject(InlineBox(parsed), forKey: key)
        return parsed
    }
}
