import Combine
import Foundation

struct AIMessage: Identifiable {
    let id: UUID
    let role: AIMessageRole
    let content: String
    let thinkingContent: String?
    let citations: [ChatCitation]?
    let isStreaming: Bool

    init(role: AIMessageRole, content: String, thinkingContent: String? = nil, citations: [ChatCitation]? = nil, isStreaming: Bool = false, id: UUID = UUID()) {
        self.id = id
        self.role = role
        self.content = content
        self.thinkingContent = thinkingContent
        self.citations = citations
        self.isStreaming = isStreaming
    }
}

struct AIMessageCitation: Identifiable {
    let id: String
    let title: String?
    let content: String?
}

enum AIMessageRole {
    case user, ai
}

@MainActor
final class AIChatViewModel: ObservableObject {
    @Published var messages: [AIMessage] = []
    @Published var inputText = ""
    @Published var isSending = false
    @Published var isCreatingSession = false
    @Published var sessionError: String?

    private var sessionId: String?
    private var entryContext: ChatEntryContext?
    private var itemIds: [String]?

    var canSend: Bool {
        !inputText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty && !isSending && sessionId != nil
    }

    var isReady: Bool { sessionId != nil }

    var scopeLabel: String? {
        guard let ctx = entryContext, ctx.scopeType != .global else { return nil }
        return "提问范围：\(ctx.scopeName)"
    }

    var scopeIcon: String {
        guard let ctx = entryContext else { return "globe" }
        switch ctx.scopeType {
        case .knowledgeBase: return "books.vertical"
        case .folder: return "folder"
        case .material: return "doc.text"
        case .knowledgeItem: return "lightbulb"
        case .global: return "globe"
        }
    }

    func load() async {
        guard sessionId == nil else { return }

        isCreatingSession = true

        let ctx = entryContext ?? ChatEntryContext(
            scopeType: .global,
            scopeId: nil,
            scopeName: "AI 对话",
            parentKnowledgeBaseId: nil,
            createdFrom: "global_ai_entry"
        )

        do {
            // open-or-create: backend handles find-or-create logic
            let session = try await RagChatService.shared.createSession(ctx: ctx)
            await loadSession(session.id)
        } catch {
            sessionError = "创建对话失败"
            isCreatingSession = false
        }
    }

    func createNewSession() async {
        sessionId = nil
        messages = []
        isCreatingSession = true

        // Use entry context if available, otherwise default to global
        let ctx = entryContext ?? ChatEntryContext(
            scopeType: .global,
            scopeId: nil,
            scopeName: "AI 对话",
            parentKnowledgeBaseId: nil,
            createdFrom: "global_ai_entry"
        )

        do {
            let session = try await RagChatService.shared.createSession(ctx: ctx)
            sessionId = session.id
            let scopeLabel = ctx.scopeType != .global ? "「\(ctx.scopeName)」的" : ""
            messages = [AIMessage(role: .ai, content: "你好！我是你的 AI 学习助手，基于\(scopeLabel)内容回答问题。")]
            isCreatingSession = false
        } catch {
            sessionError = "创建对话失败"
            isCreatingSession = false
        }
    }

    func send() {
        guard canSend, let sid = sessionId else { return }
        let text = inputText.trimmingCharacters(in: .whitespacesAndNewlines)
        messages.append(AIMessage(role: .user, content: text))
        inputText = ""
        isSending = true

        let placeholderId = UUID()
        messages.append(AIMessage(role: .ai, content: "", isStreaming: true, id: placeholderId))

        Task {
            do {
                let stream = RagChatService.shared.sendMessageStream(sessionId: sid, content: text)
                var thinking = ""
                var reply = ""

                for try await chunk in stream {
                    switch chunk.type {
                    case "thinking":
                        thinking += chunk.content ?? ""
                        if let idx = messages.firstIndex(where: { $0.id == placeholderId }) {
                            messages[idx] = AIMessage(role: .ai, content: reply, thinkingContent: thinking, isStreaming: true, id: placeholderId)
                        }
                    case "content":
                        reply += chunk.content ?? ""
                        if let idx = messages.firstIndex(where: { $0.id == placeholderId }) {
                            messages[idx] = AIMessage(role: .ai, content: reply, thinkingContent: thinking, isStreaming: true, id: placeholderId)
                        }
                    case "error":
                        if let idx = messages.firstIndex(where: { $0.id == placeholderId }) {
                            messages[idx] = AIMessage(role: .ai, content: "请求失败: \(chunk.content ?? "未知错误")")
                        }
                    case "done":
                        if let idx = messages.firstIndex(where: { $0.id == placeholderId }) {
                            messages[idx] = AIMessage(role: .ai, content: reply, thinkingContent: thinking.isEmpty ? nil : thinking)
                        }
                    default: break
                    }
                }
            } catch {
                if let idx = messages.firstIndex(where: { $0.id == placeholderId }) {
                    messages[idx] = AIMessage(role: .ai, content: "发送失败: \(error.localizedDescription)")
                }
            }
            isSending = false
        }
    }

    func setEntryContext(_ ctx: ChatEntryContext) { entryContext = ctx }

    func loadSession(_ id: String) async {
        sessionId = id
        isCreatingSession = true
        do {
            let msgs: [ChatMessage] = try await RagChatService.shared.getMessages(sessionId: id)
            if msgs.isEmpty {
                messages = [AIMessage(role: .ai, content: "你好！我是你的 AI 学习助手，基于你的知识库回答问题。")]
            } else {
                messages = msgs.map { m in
                    AIMessage(role: m.role == "user" ? .user : .ai, content: m.content, citations: m.citations)
                }
            }
        } catch { sessionError = "加载对话失败" }
        isCreatingSession = false
    }
}
