import Combine
import Foundation

struct AIMessage: Identifiable {
    let id = UUID()
    let role: AIMessageRole
    let content: String
    let citations: [ChatCitation]?
}

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 knowledgeBaseId: String?

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

    var isReady: Bool { sessionId != nil }

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

        // 自动选一个知识库
        if knowledgeBaseId == nil {
            do {
                let kbs = try await KnowledgeBaseService.shared.list(page: 1, limit: 1)
                knowledgeBaseId = kbs.first?.id
            } catch {}
        }

        guard let kbId = knowledgeBaseId else {
            sessionError = "请先创建一个知识库，才能使用 AI 对话"
            return
        }

        isCreatingSession = true
        do {
            let session = try await RagChatService.shared.createSession(knowledgeBaseId: kbId, title: nil)
            sessionId = session.id
            messages = [AIMessage(role: .ai, content: "你好！我是你的 AI 学习助手，基于你的知识库回答问题。", citations: nil)]
            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, citations: nil))
        inputText = ""
        isSending = true

        Task {
            do {
                let resp = try await RagChatService.shared.sendMessage(sessionId: sid, content: text)
                if resp.blocked == true {
                    messages.append(AIMessage(role: .ai, content: resp.message ?? "内容违规，请修改后重试", citations: nil))
                } else {
                    messages.append(AIMessage(role: .ai, content: resp.content ?? "抱歉，AI 暂时无法回复", citations: resp.citations))
                }
            } catch {
                messages.append(AIMessage(role: .ai, content: "发送失败: \(error.localizedDescription)", citations: nil))
            }
            isSending = false
        }
    }

    func setKnowledgeBase(_ id: String) { knowledgeBaseId = id }

    func loadSession(_ id: String) async {
        sessionId = id
        isCreatingSession = true
        do {
            let msgs: [ChatMessage] = try await RagChatService.shared.getMessages(sessionId: id)
            messages = msgs.map { m in
                AIMessage(role: m.role == "user" ? .user : .ai, content: m.content, citations: m.citations)
            }
        } catch { sessionError = "加载对话失败" }
        isCreatingSession = false
    }
}
