#include "localmodellistwidget.h"
#include "themedlable.h"
#include "localmodellistitem.h"
#include "dbwrapper.h"
#include "serverwrapper.h"
#include "audiocontroler.h"
#include "localmodelitem.h"
#include "localmodelserver.h"
#include "yourongmodelitem.h"
#include "modelscopeitem.h"
#include "externalllm/modelhubwrapper.h"

#include <QVBoxLayout>
#include <QHBoxLayout>

#include <DLabel>
#include <DFontSizeManager>
#include <DBackgroundGroup>
#include <DGuiApplicationHelper>

using namespace uos_ai;

LocalModelListWidget::LocalModelListWidget(DWidget *parent)
    : DWidget(parent)
{
    initUI();
    onThemeTypeChanged();
    connect(DGuiApplicationHelper::instance(), &DGuiApplicationHelper::themeTypeChanged, this, &LocalModelListWidget::onThemeTypeChanged);
}

void LocalModelListWidget::initUI()
{
    QVBoxLayout *layout = new QVBoxLayout(this);
    layout->setContentsMargins(0, 0, 0, 0);
    layout->setSpacing(0);

    m_pWidgetLabel = new ThemedLable(tr("Local model"));
    m_pWidgetLabel->setPaletteColor(QPalette::WindowText, DPalette::TextTitle);
    DFontSizeManager::instance()->bind(m_pWidgetLabel, DFontSizeManager::T6, QFont::Medium);
    layout->addWidget(m_pWidgetLabel, 0, Qt::AlignLeft);

#ifdef ENABLE_ASSISTANT
    layout->addWidget(embeddingPluginsWidget());
#ifndef COMPILE_ON_V20
    layout->addWidget(yourong1_5BLLMWidget());
    layout->addWidget(yourong7BLLMWidget());
#endif
    layout->addWidget(deepseek_1_5BLLMWidget());
#endif

    layout->addWidget(speechRecognitionWidget());
    layout->addWidget(textToImageWidget());
    layout->setSpacing(10);
    layout->addStretch();
}

void LocalModelListWidget::onThemeTypeChanged()
{
    DPalette pl = m_pTextToImageWidget->palette();
    pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
    m_pTextToImageWidget->setPalette(pl);

    pl = m_pSpeechRecognitionWidget->palette();
    pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
    m_pSpeechRecognitionWidget->setPalette(pl);

    if (m_pYouRong1_5B) {
        pl = m_pYouRong1_5B->palette();
        pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
        m_pYouRong1_5B->setPalette(pl);
    }

    if (m_pYouRong7B) {
        pl = m_pYouRong7B->palette();
        pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
        m_pYouRong7B->setPalette(pl);
    }

    if (m_pDeepseek1_5B) {
        pl = m_pDeepseek1_5B->palette();
        pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
        m_pDeepseek1_5B->setPalette(pl);
    }

    if (m_pLocalLlmWidget) {
        pl = m_pLocalLlmWidget->palette();
        pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
        m_pLocalLlmWidget->setPalette(pl);
    }

    if (m_pEmbeddingPluginsWidget) {
        pl = m_pEmbeddingPluginsWidget->palette();
        pl.setBrush(DPalette::Base, DGuiApplicationHelper::instance()->applicationPalette().color(DPalette::ItemBackground));
        m_pEmbeddingPluginsWidget->setPalette(pl);
    }
}

DBackgroundGroup *LocalModelListWidget::textToImageWidget()
{
    LocalModelListItem *item = new LocalModelListItem(this);
    item->setText(tr("Text to image model"), tr("Use a local model to generate images. After turning it on, you can select the model in the model list in the chat interface."));
    item->setFixedHeight(78);
    item->setAppName("texttoimage");
    item->setSwitchChecked(DbWrapper::localDbWrapper().queryLlmByLlmid("TextToImage").isValid());
    connect(item, &LocalModelListItem::signalSwitchChanged, this, [this](bool b) {
        updateLocalModelSwitch(LocalModel::TextToImage, b);
    });
    connect(item, &LocalModelListItem::signalUninstall, [this]() {
        updateLocalModelList();
        updateLocalModelSwitch(LocalModel::TextToImage, false);
    });

    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pTextToImageWidget = new DBackgroundGroup(bgLayout, this);
    m_pTextToImageWidget->setContentsMargins(0, 0, 0, 0);

    return m_pTextToImageWidget;
}

DBackgroundGroup *LocalModelListWidget::speechRecognitionWidget()
{
    LocalModelListItem *item = new LocalModelListItem(this);
    item->setText(tr("speech model"), tr("When turned on, the speech recognition and speech reading in the chat interface will use the local model without requiring an Internet connection."));
    item->setFixedHeight(78);
    item->setAppName("deepin-speech");
    item->setSwitchChecked(DbWrapper::localDbWrapper().getLocalSpeech());
    connect(item, &LocalModelListItem::signalSwitchChanged, this, [this](bool b) {
        updateLocalModelSwitch(LocalModel::SpeechRecognition, b);
        AudioControler::instance()->switchModel(b ? AudioControler::Local : AudioControler::NetWork);
    });
    connect(item, &LocalModelListItem::signalUninstall, [this]() {
        updateLocalModelList();
        updateLocalModelSwitch(LocalModel::SpeechRecognition, false);
        AudioControler::instance()->switchModel(AudioControler::NetWork);
    });

    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pSpeechRecognitionWidget = new DBackgroundGroup(bgLayout, this);
    m_pSpeechRecognitionWidget->setContentsMargins(0, 0, 0, 0);

    return m_pSpeechRecognitionWidget;
}

DBackgroundGroup *LocalModelListWidget::yourong1_5BLLMWidget()
{
    YouRongModelItem *item = new YouRongModelItem(this);
    item->setText(tr("ULLM-1.5B"), tr("Once installed, you do not need an internet connection to use UOS AI."));
    item->setAppName(YOURONG1_5B);
    item->setMinimumHeight(60);
    item->setMaximumHeight(75);
    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pYouRong1_5B = new DBackgroundGroup(bgLayout, this);
    m_pYouRong1_5B->setContentsMargins(0, 0, 0, 0);

    return m_pYouRong1_5B;
}

DBackgroundGroup *LocalModelListWidget::yourong7BLLMWidget()
{
    YouRongModelItem *item = new YouRongModelItem(this);
    item->setText(tr("ULLM-7B"), tr("Once installed, you do not need an internet connection to use UOS AI."));
    item->setAppName(YOURONG7B);
    item->setMinimumHeight(60);
    item->setMaximumHeight(75);
    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pYouRong7B = new DBackgroundGroup(bgLayout, this);
    m_pYouRong7B->setContentsMargins(0, 0, 0, 0);

    return m_pYouRong7B;
}

DBackgroundGroup *LocalModelListWidget::deepseek_1_5BLLMWidget()
{
    ModelScopeItem *item = new ModelScopeItem(ModelScopeItem::ItemInfo::deepseek_r1_1_5B(), this);
    item->setText(tr("DeepSeek-R1-1.5B"), tr("This model will be downloaded and installed from the open source community. Please operate with caution. After installation, you can switch to use it in UOS AI Assistant."), 1.2);
    item->setMinimumHeight(60);
    item->setMaximumHeight(75);
    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pDeepseek1_5B = new DBackgroundGroup(bgLayout, this);
    m_pDeepseek1_5B->setContentsMargins(0, 0, 0, 0);

    return m_pDeepseek1_5B;
}

DBackgroundGroup *LocalModelListWidget::localLLMWidget()
{
    LocalModelItem *item = new LocalModelItem(this);
    item->setText(tr("UOS AI large model"), tr("Once installed, you do not need an internet connection to use UOS AI."));
    item->setAppName("uos-ai-llm");
    item->setMinimumHeight(60);
    item->setMaximumHeight(75);
    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pLocalLlmWidget = new DBackgroundGroup(bgLayout, this);
    m_pLocalLlmWidget->setContentsMargins(0, 0, 0, 0);

    return m_pLocalLlmWidget;
}

DBackgroundGroup *LocalModelListWidget::embeddingPluginsWidget()
{
    LocalModelItem *item = new LocalModelItem(this);
    item->setText(tr("Embedding Plugins"), tr("After the model is installed, the System Assistant and the Personal Knowledge Assistantt can be run."));
    item->setAppName(PLUGINSNAME);
    item->setMinimumHeight(60);
    item->setMaximumHeight(75);
    QHBoxLayout *bgLayout = new QHBoxLayout;
    bgLayout->setContentsMargins(0, 0, 0, 0);
    bgLayout->addWidget(item);

    m_pEmbeddingPluginsWidget = new DBackgroundGroup(bgLayout, this);
    m_pEmbeddingPluginsWidget->setContentsMargins(0, 0, 0, 0);

    return m_pEmbeddingPluginsWidget;
}

void LocalModelListWidget::updateLocalModelList()
{
    bool textToImageExists = checkModelExist(LocalModel::TextToImage);
    bool speechRecognitionExists = checkModelExist(LocalModel::SpeechRecognition);

    m_pTextToImageWidget->setVisible(textToImageExists);
    m_pSpeechRecognitionWidget->setVisible(speechRecognitionExists);

    auto yourong1_5BWidget = m_pYouRong1_5B->findChild<YouRongModelItem *>();
    if (yourong1_5BWidget) {
        yourong1_5BWidget->checkInstallStatus();
    }
    auto yourong7BWidget = m_pYouRong7B->findChild<YouRongModelItem *>();
    if (yourong7BWidget) {
        yourong7BWidget->checkInstallStatus();
    }

    auto deepseekWidget = m_pDeepseek1_5B->findChild<ModelScopeItem *>();
    if (deepseekWidget) {
        deepseekWidget->checkInstallStatus();
    }

    LocalModelItem *llmWidget = nullptr;
    if (m_pLocalLlmWidget) {
        llmWidget = m_pLocalLlmWidget->findChild<LocalModelItem *>();
        if (llmWidget) {
            llmWidget->stopTimer();
            llmWidget->checkInstallStatus();
            llmWidget->checkUpdateStatus();
        }
    }

    auto embeddingPluginsWidget = m_pEmbeddingPluginsWidget->findChild<LocalModelItem *>();
    if (embeddingPluginsWidget) {
        embeddingPluginsWidget->stopTimer();
        embeddingPluginsWidget->checkInstallStatus();
        embeddingPluginsWidget->checkUpdateStatus();
    }
    bool isVisible = textToImageExists || speechRecognitionExists || llmWidget || embeddingPluginsWidget || yourong1_5BWidget || yourong7BWidget;

    this->setVisible(isVisible);

    this->setProperty("title", isVisible ? m_pWidgetLabel->text() : "");
}

void LocalModelListWidget::stopDownload()
{
    auto yourong1_5BWidget = m_pYouRong1_5B->findChild<YouRongModelItem *>();
    if (yourong1_5BWidget)
        yourong1_5BWidget->stopDownload();

    auto yourong7BWidget = m_pYouRong7B->findChild<YouRongModelItem *>();
    if (yourong7BWidget)
        yourong7BWidget->stopDownload();

    auto deepseekWidget = m_pDeepseek1_5B->findChild<ModelScopeItem *>();
    if (deepseekWidget) {
        deepseekWidget->stopDownload();
    }
}

bool LocalModelListWidget::checkModelExist(const LocalModel model)
{
    if (model == LocalModel::TextToImage) {
        QFileInfo binInfo("/usr/bin/texttoimage");
        return binInfo.exists() && binInfo.isFile();
    } else if (model == LocalModel::SpeechRecognition) {
        QFileInfo binInfo("/usr/bin/speechrecognition");
        return binInfo.exists() && binInfo.isFile();
    }

    return false;
}

bool LocalModelListWidget::updateLocalModelSwitch(const LocalModel model, bool b)
{
    // 已卸载，但是界面没刷新
    if (b && !checkModelExist(model)) return false;

    if (LocalModel::TextToImage == model) {
        QString modelId = "TextToImage";
        if (b) {
            LLMServerProxy newAccount;
            newAccount.name = LLMServerProxy::llmName(LLMChatModel::LOCAL_TEXT2IMAGE);
            newAccount.id = modelId;
            newAccount.model = LLMChatModel::LOCAL_TEXT2IMAGE;
            newAccount.type = ModelType::LOCAL;
            AccountProxy proxy;
            proxy.apiKey = modelId;
            newAccount.account = proxy;
            if (!DbWrapper::localDbWrapper().appendLlm(newAccount)) return false;
            ServerWrapper::instance()->updateLLMAccount();
        } else {
            if (!DbWrapper::localDbWrapper().deleteLlm(modelId)) return false;
            ServerWrapper::instance()->updateLLMAccount();
        }
    } else if (LocalModel::SpeechRecognition == model) {
        return DbWrapper::localDbWrapper().updateLocalSpeech(b);
    }

    return true;
}

QString LocalModelListWidget::getTitleName()
{
    return m_pWidgetLabel->text();
}
