import requests

class KimiTranslator:
    def __init__(self, api_key):
        self.api_key = api_key

    def translate(self, text):
        if not self.api_key:
            raise ValueError("API key is empty")
        
        url = "https://api.moonshot.ai/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        data = {
            "model": "moonshot-v1-8k",
            "messages": [
                {"role": "system", "content": "You are a helpful and precise assistant for translating Chinese to English."},
                {"role": "user", "content": f"Translate the following Chinese text to English: {text}"}
            ],
            "temperature": 0.3
        }
        
        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            return response.json()["choices"][0]["message"]["content"]
        else:
            raise Exception(f"Error: {response.status_code} - {response.text}")

import sys
from PyQt5 import QtWidgets
import cv2
from PyQt5.QtGui import *
from PyQt5.QtCore import *
import numpy as np
from PyQt5.QtWidgets import *
from ultralytics import solutions
import os
from main_window import Ui_MainWindow
from KimiTranslator import KimiTranslator

class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.dir = "JPEGImages"
        self.searcher = solutions.VisualAISearch(
            data=self.dir
        )
        # 填寫 Kimi 的 API Key
        API_KEY = ""  # 在這里替換為你的 API Key
        
        # 創(chuàng)建翻譯器實(shí)例
        self.translator = KimiTranslator(API_KEY)
        
        # 初始化文件路徑
        self.file_path = None
        
        # 連接按鈕信號(hào)（假設(shè) UI 中的對(duì)象名；如果不同，請(qǐng)調(diào)整）
        self.btnSelectImage.clicked.connect(self.selectImage)  # 選擇圖像按鈕
        self.btnTextSearch.clicked.connect(self.searchText2Pic)  # 以文搜圖按鈕
        self.btnImageSearch.clicked.connect(self.searchPic2Pic)  # 以圖搜圖按鈕

    def updateImage(self, path, qlabel):
        image = self.cv_imread(path)
        if image is not None:
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            rows, cols, channels = rgb_image.shape
            bytesPerLine = channels * cols
            QImg = QImage(rgb_image.data, cols, rows, bytesPerLine, QImage.Format_RGB888)
            qlabel.setPixmap(QPixmap.fromImage(QImg).scaled(qlabel.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def selectImage(self):
        file_path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Select an Image", "", "Image Files (*.jpg *.jpeg *.png)")
        if file_path:
            self.file_path = file_path
            self.updateImage(self.file_path, self.queryImagePreview)

    def cv_imread(self, path):
        cv_img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
        return cv_img

    def searchText2Pic(self):
        query_text = self.textQueryEdit.toPlainText().strip()
        if not query_text:
            QtWidgets.QMessageBox.warning(self, "提示", "輸入文本不能為空")
            return
        translated = self.translator.translate(query_text)
        results = self.searcher.search_text(translated)
        self.updateResults(results)

    def searchPic2Pic(self):
        if not self.file_path:
            QtWidgets.QMessageBox.warning(self, "提示", "請(qǐng)選擇圖像")
            return
        results = self.searcher.search_image(self.file_path)
        self.updateResults(results)

    def updateResults(self, results):
        image_paths = []
        similarities = []
        for image_path, similarity in results[:8]:
            if hasattr(image_path, '__str__'):
                image_path = str(image_path)
            # 拼接完整路徑
            full_path = os.path.join(self.dir, image_path)
            image_paths.append(full_path)
            similarities.append(round(similarity, 4))
        
        # 更新圖像和相似度，處理少于8個(gè)結(jié)果的情況
        for i in range(8):
            label = getattr(self, f'image_{i+1}')
            sim_label = getattr(self, f'similarity_{i+1}')
            if i < len(image_paths):
                self.updateImage(image_paths[i], label)
                sim_label.setText("相似度：" + str(similarities[i]))
            else:
                label.clear()
                sim_label.setText("")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())