import sys
import requests
import json
from PySide6.QtWidgets import QApplication, QMainWindow, QListWidget, QVBoxLayout, QWidget, QLabel, QPushButton
from PySide6.QtCore import Qt
from loguru import logger

class OllamaModelViewer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Ollama 模型查看器")
        self.setGeometry(100, 100, 600, 400)
        
        # 创建主窗口部件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        
        # 创建布局
        self.layout = QVBoxLayout()
        self.central_widget.setLayout(self.layout)
        
        # 创建标题标签
        self.title_label = QLabel("当前安装的Ollama模型:")
        self.title_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        self.layout.addWidget(self.title_label)
        
        # 创建列表部件
        self.model_list = QListWidget()
        self.model_list.setStyleSheet("font-family: monospace;")
        self.layout.addWidget(self.model_list)
        
        # 创建刷新按钮
        self.refresh_button = QPushButton("刷新模型列表")
        self.refresh_button.clicked.connect(self.fetch_models)
        self.layout.addWidget(self.refresh_button)
        
        # 初始加载模型
        self.fetch_models()
    
    def fetch_models(self):
        """从Ollama API获取模型列表"""
        self.model_list.clear()
        
        try:
            logger.info("正在获取Ollama模型列表...")
            response = requests.get("http://localhost:11434/api/tags", timeout=5)
            
            if response.status_code == 200:
                data = response.json()
                models = data.get("models", [])
                
                if models:
                    for model in models:
                        model_name = model.get("model", "")
                        if model_name:
                            self.model_list.addItem(model_name)
                            logger.info(f"找到模型: {model_name}")
                else:
                    self.model_list.addItem("未找到任何模型")
                    logger.info("未找到任何模型")
            else:
                self.model_list.addItem(f"API请求失败，状态码: {response.status_code}")
                logger.error(f"API请求失败，状态码: {response.status_code}")
                
        except requests.exceptions.RequestException as e:
            self.model_list.addItem("无法连接到Ollama API")
            logger.error(f"无法连接到Ollama API: {str(e)}")
        except json.JSONDecodeError as e:
            self.model_list.addItem("API响应格式错误")
            logger.error(f"API响应格式错误: {str(e)}")
        except Exception as e:
            self.model_list.addItem(f"发生错误: {str(e)}")
            logger.error(f"发生未知错误: {str(e)}")

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