import Foundation
import NucleicCore
import OSLog
import SwiftUI

struct MarkdownPrewarmStats: Equatable, Sendable {
    var documents = 0
    var memoryHits = 0
    var diskHits = 0
    var parsed = 0
    var elapsedMilliseconds: Double = 0
}

enum MarkdownPrewarmDiskAccess: Sendable, Equatable {
    /// Read the session snapshot and replace it with the exact supplied document set.
    case sessionSnapshot
    /// Populate only the process cache. Used by long-history first paint so decoding or rewriting
    /// a large session sidecar can never sit between selection and visible text.
    case memoryOnly
}

/// 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: Sendable { 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: @unchecked Sendable {
        let blocks: [Block]
        init(_ b: [Block]) { blocks = b }
    }
    nonisolated(unsafe) private static let blockCache: NSCache<NSString, ParsedBlocks> = {
        let cache = NSCache<NSString, ParsedBlocks>()
        cache.countLimit = 2048
        return cache
    }()

    nonisolated 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
    }

    nonisolated 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.
    nonisolated 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])
    }

    nonisolated 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.
    nonisolated 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))
        // A fenced block is the agent's way of handing over something copy-verbatim
        // (a command, a snippet), so give it a one-click copy on hover.
        .copyable(code, help: "Copy code")
    }

    /// A run of prose lines rendered as a *single* `Text`, not one `Text` per line.
    /// SwiftUI text selection can't span sibling `Text` views, so the old per-line
    /// `VStack` made it impossible to select/copy across a line break — the whole
    /// block is now one selectable unit, with per-line styling carried by the runs.
    private func textBlock(_ text: String) -> some View {
        let lines = text.components(separatedBy: "\n")
        var composed = Text(verbatim: "")
        for (idx, raw) in lines.enumerated() {
            if idx > 0 { composed = Text("\(composed)\n") }
            composed = Text("\(composed)\(lineText(raw))")
        }
        return composed
            .frame(maxWidth: .infinity, alignment: .leading)
            .multilineTextAlignment(.leading)
            // Breathing room between the visual lines of a single wrapped paragraph.
            .lineSpacing(4)
            .textSelection(.enabled)
    }

    /// Style one source line into a `Text` run. Returns styled runs (heading fonts,
    /// bullet glyphs) that concatenate into the block's single selectable `Text`.
    private func lineText(_ raw: String) -> Text {
        let trimmed = raw.trimmingCharacters(in: .whitespaces)
        if trimmed.isEmpty {
            // The surrounding "\n" joins already supply the blank line's break.
            return Text(verbatim: "")
        } 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.
            return inline(String(trimmed.dropFirst(4))).font(.system(size: bodySize * 1.13, weight: .semibold))
        } else if trimmed.hasPrefix("## ") {
            return inline(String(trimmed.dropFirst(3))).font(.system(size: bodySize * 1.28, weight: .bold))
        } else if trimmed.hasPrefix("# ") {
            return inline(String(trimmed.dropFirst(2))).font(.system(size: bodySize * 1.5, weight: .bold))
        } else if let bullet = Self.bulletContent(trimmed) {
            // Inline the marker into the run (vs. a separate HStack column) so a wrapped
            // bullet stays part of the one selectable block.
            return Text("\(Text("•  ").foregroundStyle(.secondary))\(inline(bullet))")
        } else {
            return inline(raw)
        }
    }

    /// Returns the content after a `- `, `* `, `+ ` or `N. ` list marker, else nil.
    nonisolated 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: @unchecked Sendable {
        let value: AttributedString
        init(_ v: AttributedString) { value = v }
    }
    nonisolated(unsafe) private static let inlineCache: NSCache<NSString, InlineBox> = {
        let cache = NSCache<NSString, InlineBox>()
        cache.countLimit = 16384
        return cache
    }()

    nonisolated private static func attributedInline(_ string: String) -> AttributedString {
        let key = string as NSString
        if let hit = inlineCache.object(forKey: key) { return hit.value }
        let parsed = parseInlineUncached(string)
        inlineCache.setObject(InlineBox(parsed), forKey: key)
        return parsed
    }

    nonisolated private static func parseInlineUncached(_ string: String) -> AttributedString {
        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.
        return (try? AttributedString(markdown: string, options: options)) ?? AttributedString(string)
    }

    // MARK: - Cold-open prewarming + durable parse cache

    /// Load a session's parsed inline Markdown from disk, parse only misses off-main, and populate
    /// the same `NSCache` the SwiftUI body reads synchronously. A long transcript first calls this
    /// for its visible tail in `.memoryOnly` mode, reveals the virtualized list, then loads/refreshes
    /// the complete session snapshot in the background. Short eager transcripts still warm their
    /// complete history before first layout.
    ///
    /// The disk representation is renderer-neutral: text runs + presentation intents + links.
    /// Width, fonts, colors, and SwiftUI layout remain live inputs and are never persisted.
    nonisolated static func prewarm(
        _ markdowns: [String],
        sessionID: SessionID,
        diskCache: MarkdownRenderDiskCache = .shared,
        diskAccess: MarkdownPrewarmDiskAccess = .sessionSnapshot
    ) async -> MarkdownPrewarmStats {
        let started = Date()
        let signposter = OSSignposter(subsystem: "com.nucleic", category: "transcript-markdown")
        let signpostID = signposter.makeSignpostID()
        let interval = signposter.beginInterval("MarkdownPrewarm", id: signpostID)
        defer { signposter.endInterval("MarkdownPrewarm", interval) }

        var stats = MarkdownPrewarmStats()
        let stored: [String: MarkdownCachedDocument]
        switch diskAccess {
        case .sessionSnapshot:
            stored = await diskCache.load(sessionID: sessionID)
        case .memoryOnly:
            stored = [:]
        }
        var current: [String: MarkdownCachedDocument] = [:]

        // A repeated message body is one parse/cache entry. Preserve order so the tail (the rows
        // first exposed on open) warms first when cancellation races a session switch.
        var seen = Set<String>()
        let unique = markdowns.reversed().filter { seen.insert($0).inserted }
        for markdown in unique {
            if Task.isCancelled { break }
            stats.documents += 1
            let documentKey = MarkdownRenderDiskCache.documentKey(for: markdown)
            let storedDocument = stored[documentKey]
            let storedInlines = Dictionary(
                uniqueKeysWithValues: (storedDocument?.inlines ?? []).map { ($0.source, $0) })
            var cachedInlines: [MarkdownCachedInline] = []

            for source in inlineSources(in: markdown) {
                if Task.isCancelled { break }
                let key = source as NSString
                if let memory = inlineCache.object(forKey: key)?.value {
                    stats.memoryHits += 1
                    cachedInlines.append(cacheRecord(source: source, attributed: memory))
                } else if storedDocument?.sourceByteCount == markdown.utf8.count,
                          let disk = storedInlines[source]
                {
                    let attributed = attributedString(from: disk)
                    inlineCache.setObject(InlineBox(attributed), forKey: key)
                    stats.diskHits += 1
                    cachedInlines.append(disk)
                } else {
                    let parsed = parseInlineUncached(source)
                    inlineCache.setObject(InlineBox(parsed), forKey: key)
                    stats.parsed += 1
                    cachedInlines.append(cacheRecord(source: source, attributed: parsed))
                }
            }
            current[documentKey] = MarkdownCachedDocument(
                sourceByteCount: markdown.utf8.count,
                inlines: cachedInlines)
        }

        if !Task.isCancelled, diskAccess == .sessionSnapshot, current != stored {
            await diskCache.save(current, sessionID: sessionID)
        }
        stats.elapsedMilliseconds = Date().timeIntervalSince(started) * 1_000
        let logger = Logger(subsystem: "com.nucleic", category: "transcript-markdown")
        if diskAccess == .sessionSnapshot {
            logger.info(
                "prewarm documents=\(stats.documents) memory=\(stats.memoryHits) disk=\(stats.diskHits) parsed=\(stats.parsed) elapsed_ms=\(stats.elapsedMilliseconds, format: .fixed(precision: 1))")
        } else {
            // Memory-only work includes the continuous long-history pass and viewport-priority
            // fills. Keep its measurement available without flooding the default log.
            logger.debug(
                "segment_prewarm documents=\(stats.documents) memory=\(stats.memoryHits) parsed=\(stats.parsed) elapsed_ms=\(stats.elapsedMilliseconds, format: .fixed(precision: 1))")
        }
        return stats
    }

    /// The exact strings that flow into `AttributedString(markdown:)` during rendering.
    nonisolated private static func inlineSources(in markdown: String) -> [String] {
        var result: [String] = []
        var seen = Set<String>()
        func append(_ source: String) {
            if seen.insert(source).inserted { result.append(source) }
        }
        for block in parse(markdown) {
            switch block {
            case .code:
                break
            case .table(let rows):
                for row in rows { for cell in row { append(cell) } }
            case .text(let text):
                for raw in text.components(separatedBy: "\n") {
                    if let source = inlineSource(for: raw) { append(source) }
                }
            }
        }
        return result
    }

    nonisolated private static func inlineSource(for raw: String) -> String? {
        let trimmed = raw.trimmingCharacters(in: .whitespaces)
        guard !trimmed.isEmpty else { return nil }
        if trimmed.hasPrefix("### ") { return String(trimmed.dropFirst(4)) }
        if trimmed.hasPrefix("## ") { return String(trimmed.dropFirst(3)) }
        if trimmed.hasPrefix("# ") { return String(trimmed.dropFirst(2)) }
        if let bullet = bulletContent(trimmed) { return bullet }
        return raw
    }

    nonisolated private static func cacheRecord(
        source: String, attributed: AttributedString
    ) -> MarkdownCachedInline {
        var runs: [MarkdownCachedRun] = []
        for run in attributed.runs {
            let record = MarkdownCachedRun(
                text: String(attributed.characters[run.range]),
                presentationIntent: run.inlinePresentationIntent?.rawValue,
                link: run.link?.absoluteString)
            runs.append(record)
        }
        return MarkdownCachedInline(source: source, runs: runs)
    }

    nonisolated private static func attributedString(
        from cached: MarkdownCachedInline
    ) -> AttributedString {
        var attributed = AttributedString()
        for record in cached.runs {
            var run = AttributedString(record.text)
            if let rawValue = record.presentationIntent {
                run.inlinePresentationIntent = InlinePresentationIntent(rawValue: rawValue)
            }
            if let link = record.link, let url = URL(string: link) { run.link = url }
            attributed.append(run)
        }
        return attributed
    }

    #if DEBUG
    nonisolated static func clearMemoryCachesForTesting() {
        blockCache.removeAllObjects()
        inlineCache.removeAllObjects()
    }

    nonisolated static func cachedInlineForTesting(_ source: String) -> AttributedString? {
        inlineCache.object(forKey: source as NSString)?.value
    }
    #endif
}
