from typing import Dict, List, Optional

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import aiohttp
import asyncio
import json
import sys
import pkg_resources
import subprocess
from typing import Dict, List, Optional

class GeminiChat:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://gemini.tangrenbin.us.kg"
        self.messages = []
        self.conversation_history = []  # 存储对话历史
        self.max_history = 10  # 默认保留最近10轮对话

    @staticmethod
    def check_dependencies() -> None:
        """检查并安装必要的依赖"""
        required = {'aiohttp', 'asyncio'}
        installed = {pkg.key for pkg in pkg_resources.working_set}
        missing = required - installed

        if missing:
            print("检测到缺少必要的依赖项，正在安装...")
            try:
                subprocess.check_call([sys.executable, '-m', 'pip', 'install', *missing])
                print("依赖安装完成！")
            except subprocess.CalledProcessError:
                print("错误：依赖安装失败，请手动安装以下包：")
                print("\n".join(missing))
                sys.exit(1)

    async def fetch_models(self) -> Dict[str, str]:
        """从API获取可用的模型列表"""
        print("正在获取可用模型列表...")
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    f"{self.base_url}/v1/models",
                    headers={'Authorization': f'Bearer {self.api_key}'}
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        models = {}
                        index = 1
                        for model in data['data']:
                            if 'gemini' in model['id']:
                                models[str(index)] = model['id']
                                index += 1
                        return models
                    else:
                        error_text = await response.text()
                        raise Exception(f"获取模型列表失败: {error_text}")
        except Exception as e:
            print(f"错误：无法获取模型列表 - {str(e)}")
            sys.exit(1)

    def print_models(self, models: Dict[str, str], default_model: str) -> None:
        """打印模型列表"""
        print("\n可用模型列表：")
        print("注意：不同模型有不同的特点：")
        print("- pro系列：响应质量更高")
        print("- flash系列：响应速度更快")
        print("- 数字越大代表版本越新")
        print(f"- 默认使用最新的模型: {models[default_model]}\n")

        for key, model in models.items():
            print(f"{key}. {model}" + (" (默认)" if key == default_model else ""))

    def save_conversation(self, role: str, content: str) -> None:
        """保存对话历史"""
        self.conversation_history.append({"role": role, "content": content})
        if len(self.conversation_history) > self.max_history * 2:  # 每轮对话包含用户和AI两条消息
            self.conversation_history = self.conversation_history[-self.max_history * 2:]

    def get_conversation_context(self) -> List[Dict[str, str]]:
        """获取对话上下文"""
        return [{"role": "system", "content": "你是一个有帮助的AI助手。"}] + self.conversation_history

    async def chat_completion(self, model: str, temperature: float = 0.7) -> str:
        """发送聊天请求"""
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self.base_url}/v1/chat/completions",
                    headers={
                        'Content-Type': 'application/json',
                        'Authorization': f'Bearer {self.api_key}'
                    },
                    json={
                        "model": model,
                        "messages": self.get_conversation_context(),
                        "stream": False,
                        "temperature": temperature,
                        "max_tokens": 2048
                    }
                ) as response:
                    if response.status == 200:
                        data = await response.json()
                        return data['choices'][0]['message']['content']
                    else:
                        error_text = await response.text()
                        raise Exception(f"API错误: {error_text}")
        except Exception as e:
            return f"错误：{str(e)}"

    async def chat(self) -> None:
        """主对话循环"""
        # 获取模型列表
        models = await self.fetch_models()
        default_model = str(len(models))  # 使用最后一个模型作为默认值
        
        while True:
            self.print_models(models, default_model)
            print(f"\n直接按回车使用默认模型({models[default_model]})")
            model_choice = input("请选择模型编号: ").strip()
            
            if not model_choice:
                model_choice = default_model
                
            if model_choice in models:
                break
            print("无效的选择，请重试。")
        
        model = models[model_choice]
        print(f"\n已选择模型: {model}")
        print("\n开始对话 (输入'退出'结束对话，输入'切换'可以切换模型，输入'刷新'可以刷新模型列表，输入'清空'可以清空对话历史)")
        
        while True:
            user_input = input("\n你: ").strip()
            
            if user_input.lower() in ['退出', 'quit', 'exit']:
                print("\n感谢使用！再见！")
                break
                
            if user_input.lower() == '切换':
                self.conversation_history = []  # 切换模型时清空对话历史
                await self.chat()
                return
                
            if user_input.lower() == '刷新':
                models = await self.fetch_models()
                self.print_models(models, model_choice)
                print(f"\n继续使用当前模型: {model}")
                continue
                
            if user_input.lower() == '清空':
                self.conversation_history = []
                print("\n对话历史已清空")
                continue
            
            if user_input:
                # 保存用户输入
                self.save_conversation("user", user_input)
                
                print("\n正在等待AI响应...")
                response = await self.chat_completion(model)
                print(f"\nAI: {response}")
                
                # 保存AI回复
                self.save_conversation("assistant", response)

async def main():
    # 检查依赖
    GeminiChat.check_dependencies()
    
    api_key = "AIzaSyDOLKXvMTqw0EgLHVOA13fgBZHbS8NEvBs"
    chat = GeminiChat(api_key)
    await chat.chat()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n程序已被用户中断。再见！")
    except Exception as e:
        print(f"\n程序遇到错误: {str(e)}")