import SwiftUI

/// Session control pills (model + thinking).
///
/// Body returns flat children (no wrapping container) so they merge
/// into the parent's `HStack` for even distribution with the attach button.
struct SessionToolbar: View {
    let session: Session?
    var modelOverride: String? = nil
    let thinkingLevel: ThinkingLevel
    let onModelTap: () -> Void
    let onThinkingSelect: (ThinkingLevel) -> Void

    @Environment(\.theme) private var theme
    @State private var isThinkingPickerPresented = false

    private var effectiveModel: String? {
        modelOverride ?? session?.model
    }

    private var modelDisplay: String {
        guard let model = effectiveModel else { return "default" }
        return shortModelName(model)
    }

    private var modelProvider: String? {
        guard let model = effectiveModel else { return nil }
        return providerFromModel(model)
    }

    private var thinkingTint: Color {
        theme.thinking.color(for: thinkingLevel)
    }

    // A system Menu can reverse its visual order when it expands upward from a
    // keyboard-adjacent composer. The fixed popover stack keeps this order stable.
    private var thinkingLabel: String {
        thinkingLevel.compactTitle
    }

    var body: some View {
        Spacer(minLength: 0)

        Button(action: onModelTap) {
            PillLabel(
                text: modelDisplay,
                showChevron: true,
                showsLeadingIcon: modelProvider != nil
            ) {
                if let modelProvider {
                    ProviderIcon(provider: modelProvider)
                }
            }
        }
        .buttonStyle(.plain)
        .accessibilityIdentifier("session.toolbar.model")

        Button {
            isThinkingPickerPresented = true
        } label: {
            PillLabel(text: thinkingLabel, tint: thinkingTint, showChevron: true) {
                Image(systemName: "sparkle")
                    .font(.appCaption)
                    .symbolRenderingMode(.hierarchical)
            }
        }
        .buttonStyle(.plain)
        .popover(isPresented: $isThinkingPickerPresented, arrowEdge: .bottom) {
            VStack(spacing: 0) {
                ForEach(ThinkingLevel.allCases.reversed(), id: \.rawValue) { level in
                    Button {
                        isThinkingPickerPresented = false
                        guard level != thinkingLevel else { return }
                        onThinkingSelect(level)
                    } label: {
                        HStack(spacing: 10) {
                            Text(level.displayTitle)
                                .frame(maxWidth: .infinity, alignment: .leading)

                            Image(systemName: "checkmark")
                                .opacity(level == thinkingLevel ? 1 : 0)
                                .accessibilityHidden(true)
                        }
                        .font(.body)
                        .foregroundStyle(.themeFg)
                        .padding(.horizontal, 16)
                        .frame(minHeight: 44)
                        .contentShape(Rectangle())
                    }
                    .buttonStyle(.plain)
                    .accessibilityLabel(level.displayTitle)
                    .accessibilityValue(level == thinkingLevel ? "Selected" : "")
                    .accessibilityIdentifier("session.toolbar.thinking.option.\(level.rawValue)")
                }
            }
            .padding(.vertical, 6)
            .frame(width: 190)
            .presentationCompactAdaptation(.popover)
            .presentationBackground(Color.themeSurfaceFill(.popover))
        }
        .accessibilityLabel("Thinking level")
        .accessibilityValue(thinkingLevel.displayTitle)
        .accessibilityIdentifier("session.toolbar.thinking")
    }
}

private struct PillLabel<LeadingIcon: View>: View {
    private let iconSlotSize: CGFloat = 11
    private let chevronSlotWidth: CGFloat = 9
    private let minPillHeight: CGFloat = 17

    let text: String
    var tint: Color
    let showChevron: Bool
    let showsLeadingIcon: Bool
    let leadingIcon: LeadingIcon

    init(
        text: String,
        tint: Color = .themeFg,
        showChevron: Bool,
        showsLeadingIcon: Bool = true,
        @ViewBuilder leadingIcon: () -> LeadingIcon
    ) {
        self.text = text
        self.tint = tint
        self.showChevron = showChevron
        self.showsLeadingIcon = showsLeadingIcon
        self.leadingIcon = leadingIcon()
    }

    var body: some View {
        HStack(alignment: .center, spacing: 3) {
            if showsLeadingIcon {
                leadingIcon
                    .frame(
                        width: iconSlotSize,
                        height: iconSlotSize,
                        alignment: .center
                    )
            }

            Text(text)
                .font(.caption2.monospacedDigit().weight(.medium))
                .lineLimit(1)
                .truncationMode(.tail)

            if showChevron {
                Image(systemName: "chevron.down")
                    .font(.appBadgeCount)
                    .frame(
                        width: chevronSlotWidth,
                        height: iconSlotSize,
                        alignment: .center
                    )
            }
        }
        .frame(minHeight: minPillHeight)
        .foregroundStyle(tint)
        .padding(.horizontal, 8)
        .padding(.vertical, 2)
        .glassEffect(.regular, in: Capsule())
        .contentShape(Capsule())
    }
}

/// Best-effort context window fallback when older sessions lack the field.
/// Static lookup table for known model context windows.
/// Allocated once, shared across all calls.
private let knownContextWindows: [String: Int] = [
    "anthropic/claude-opus-4-6": 200_000,
    "anthropic/claude-sonnet-4-0": 200_000,
    "anthropic/claude-haiku-3-5": 200_000,
    "openai/o3": 200_000,
    "openai/o4-mini": 200_000,
    "openai/gpt-4.1": 1_000_000,
    "openai-codex/gpt-5.1": 272_000,
    "openai-codex/gpt-5.2": 272_000,
    "openai-codex/gpt-5.2-codex": 272_000,
    "openai-codex/gpt-5.3-codex": 272_000,
    "google/gemini-2.5-pro": 1_000_000,
    "google/gemini-2.5-flash": 1_000_000,
    "lmstudio/glm-4.7": 128_000,
    "lmstudio/glm-4.7-flash-mlx": 128_000,
    "lmstudio/magistral-small-2509-mlx": 32_000,
    "lmstudio/minimax-m2.1": 196_608,
    "lmstudio/qwen3-coder-next": 128_000,
    "lmstudio/qwen3-32b": 32_768,
    "lmstudio/deepseek-r1-0528-qwen3-8b": 32_768,
]

func inferContextWindow(from model: String) -> Int? {
    if let value = knownContextWindows[model] {
        return value
    }

    // Fast manual scan for "NNNk" or "NNNk" pattern (2-4 digits followed by k/K at word boundary).
    // Avoids regex compilation and NSString bridging.
    let utf8 = Array(model.utf8)
    let count = utf8.count
    var i = count - 1
    // Scan backwards — the context size suffix is typically at the end
    while i >= 1 {
        let c = utf8[i]
        if c == UInt8(ascii: "k") || c == UInt8(ascii: "K") {
            // Check that next char (if any) is not alphanumeric (word boundary)
            let afterK = i + 1
            if afterK < count {
                let next = utf8[afterK]
                let isAlnum = (next >= UInt8(ascii: "0") && next <= UInt8(ascii: "9"))
                    || (next >= UInt8(ascii: "a") && next <= UInt8(ascii: "z"))
                    || (next >= UInt8(ascii: "A") && next <= UInt8(ascii: "Z"))
                    || next == UInt8(ascii: "_")
                if isAlnum { i -= 1; continue }
            }
            // Collect 2-4 digits before k/K
            let digitEnd = i
            var j = i - 1
            while j >= 0 && utf8[j] >= UInt8(ascii: "0") && utf8[j] <= UInt8(ascii: "9") {
                j -= 1
            }
            let digitStart = j + 1
            let digitCount = digitEnd - digitStart
            if digitCount >= 2 && digitCount <= 4 {
                var value = 0
                for d in digitStart..<digitEnd {
                    value = value * 10 + Int(utf8[d] - UInt8(ascii: "0"))
                }
                return value * 1_000
            }
        }
        i -= 1
    }

    return nil
}

/// Extract provider prefix from "provider/model-id" string.
func providerFromModel(_ model: String) -> String? {
    guard let slashIndex = model.firstIndex(of: "/") else { return nil }
    let provider = String(model[model.startIndex..<slashIndex])
    return provider.isEmpty ? nil : provider
}

/// Extract short display name from full "provider/model-id" string.
func shortModelName(_ model: String) -> String {
    let name = model.split(separator: "/").last.map(String.init) ?? model
    return name
        .replacingOccurrences(of: "claude-", with: "")
        .replacingOccurrences(of: "gemini-", with: "")
}
