#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Llama API 通用對話程式
支援內網和外網端點選擇
"""

from openai import OpenAI
import requests
import sys
import re
from datetime import datetime

# API 金鑰
API_KEY = "paVrIT+XU1NhwCAOb0X4aYi75QKogK5YNMGvQF1dCyo="

# 網路環境配置
NETWORK_CONFIG = {
    "內網": {
        "name": "內網環境",
        "description": "公司/學校內部網路，使用內部 IP 地址",
        "endpoints": [
            {
                "name": "內網端點 1",
                "url": "http://192.168.0.6:21180/v1",
                "models": ["gpt-oss-120b", "deepseek-r1-671b", "qwen3-embedding-8b"]
            },
            {
                "name": "內網端點 2",
                "url": "http://192.168.0.6:21181/v1",
                "models": ["gpt-oss-120b", "deepseek-r1-671b", "qwen3-embedding-8b"]
            },
            {
                "name": "內網端點 3",
                "url": "http://192.168.0.6:21182/v1",
                "models": ["gpt-oss-120b", "deepseek-r1-671b", "qwen3-embedding-8b"]
            }
        ]
    },
    "外網": {
        "name": "外網環境",
        "description": "公開網際網路，使用外部域名",
        "endpoints": [
            {
                "name": "通用端點",
                "url": "https://llama.theaken.com/v1",
                "models": ["gpt-oss-120b", "deepseek-r1-671b", "qwen3-embedding-8b"]
            },
            {
                "name": "GPT-OSS 專用端點",
                "url": "https://llama.theaken.com/v1/gpt-oss-120b",
                "models": ["gpt-oss-120b"]
            },
            {
                "name": "DeepSeek 專用端點",
                "url": "https://llama.theaken.com/v1/deepseek-r1-671b",
                "models": ["deepseek-r1-671b"]
            }
        ]
    }
}

def clean_response(text):
    """清理 AI 回應中的特殊標記"""
    if not text:
        return text
        
    # 移除思考標記
    if "<think>" in text:
        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
    
    # 移除 channel 標記
    if "<|channel|>" in text:
        parts = text.split("<|message|>")
        if len(parts) > 1:
            text = parts[-1]
    
    # 移除結束標記
    text = text.replace("<|end|>", "").replace("<|start|>", "")
    
    # 清理多餘空白
    text = text.strip()
    
    return text

def test_endpoint(endpoint_info, timeout=8):
    """測試端點是否可用"""
    url = endpoint_info["url"]
    model = endpoint_info["models"][0] if endpoint_info["models"] else "gpt-oss-120b"
    
    print(f"  測試 {endpoint_info['name']}...", end="", flush=True)
    
    try:
        # 處理特殊的模型端點 URL
        if url.endswith("/gpt-oss-120b") or url.endswith("/deepseek-r1-671b"):
            base_url = url.rsplit("/", 1)[0]
        else:
            base_url = url
            
        client = OpenAI(
            api_key=API_KEY, 
            base_url=base_url,
            timeout=timeout
        )
        
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": "test"}],
            max_tokens=5
        )
        print(" [OK]")
        return True
        
    except Exception as e:
        error_msg = str(e)
        if "502" in error_msg:
            print(" [502 - 伺服器離線]")
        elif "timeout" in error_msg.lower():
            print(" [超時]")
        elif "connection" in error_msg.lower():
            print(" [連接失敗]")
        else:
            print(" [錯誤]")
        return False

def choose_network_environment():
    """選擇網路環境"""
    print("\n" + "="*60)
    print("選擇網路環境")
    print("="*60)
    
    print("\n可用環境：")
    environments = list(NETWORK_CONFIG.keys())
    
    for i, env in enumerate(environments, 1):
        config = NETWORK_CONFIG[env]
        print(f"  {i}. {config['name']}")
        print(f"     說明：{config['description']}")
        print(f"     端點數量：{len(config['endpoints'])} 個")
    
    while True:
        try:
            choice = input(f"\n請選擇環境 (1-{len(environments)})，預設為 1: ").strip()
            
            if not choice:
                choice = "1"
                
            choice_num = int(choice)
            if 1 <= choice_num <= len(environments):
                selected_env = environments[choice_num - 1]
                print(f"\n已選擇：{NETWORK_CONFIG[selected_env]['name']}")
                return selected_env
            else:
                print(f"請輸入 1-{len(environments)} 之間的數字")
                
        except ValueError:
            print("請輸入有效的數字")
        except KeyboardInterrupt:
            print("\n\n程式已取消")
            sys.exit(0)

def find_working_endpoint(environment):
    """尋找指定環境中的可用端點"""
    config = NETWORK_CONFIG[environment]
    
    print(f"\n正在測試 {config['name']} 的端點...")
    print("-" * 50)
    
    for endpoint in config['endpoints']:
        if test_endpoint(endpoint):
            return endpoint
    
    return None

def chat_session(endpoint_info, environment):
    """對話主程式"""
    print("\n" + "="*60)
    print("Llama AI 對話系統")
    print("="*60)
    print(f"網路環境: {NETWORK_CONFIG[environment]['name']}")
    print(f"使用端點: {endpoint_info['name']}")
    print(f"URL: {endpoint_info['url']}")
    print(f"可用模型: {', '.join(endpoint_info['models'])}")
    print("\n指令:")
    print("  exit/quit - 結束對話")
    print("  clear - 清空對話歷史")
    print("  model - 切換模型")
    print("  switch - 切換網路環境")
    print("-"*60)
    
    # 處理 URL
    url = endpoint_info["url"]
    if url.endswith("/gpt-oss-120b") or url.endswith("/deepseek-r1-671b"):
        base_url = url.rsplit("/", 1)[0]
    else:
        base_url = url
    
    client = OpenAI(api_key=API_KEY, base_url=base_url)
    
    # 選擇模型
    if len(endpoint_info['models']) == 1:
        current_model = endpoint_info['models'][0]
    else:
        print("\n選擇模型:")
        for i, model in enumerate(endpoint_info['models'], 1):
            print(f"  {i}. {model}")
        choice = input("選擇 (預設: 1): ").strip()
        if choice.isdigit() and 1 <= int(choice) <= len(endpoint_info['models']):
            current_model = endpoint_info['models'][int(choice)-1]
        else:
            current_model = endpoint_info['models'][0]
    
    print(f"\n使用模型: {current_model}")
    messages = []
    
    while True:
        try:
            user_input = input("\n你: ").strip()
            
            if not user_input:
                continue
                
            if user_input.lower() in ['exit', 'quit']:
                print("再見！")
                break
                
            if user_input.lower() == 'clear':
                messages = []
                print("[系統] 對話歷史已清空")
                continue
                
            if user_input.lower() == 'switch':
                print("[系統] 重新啟動程式以切換網路環境")
                return "switch"
                
            if user_input.lower() == 'model':
                if len(endpoint_info['models']) == 1:
                    print(f"[系統] 此端點只支援 {endpoint_info['models'][0]}")
                else:
                    print("\n可用模型:")
                    for i, m in enumerate(endpoint_info['models'], 1):
                        print(f"  {i}. {m}")
                    choice = input("選擇: ").strip()
                    if choice.isdigit() and 1 <= int(choice) <= len(endpoint_info['models']):
                        current_model = endpoint_info['models'][int(choice)-1]
                        print(f"[系統] 已切換到 {current_model}")
                continue
            
            messages.append({"role": "user", "content": user_input})
            
            print("\nAI 思考中...", end="", flush=True)
            
            try:
                response = client.chat.completions.create(
                    model=current_model,
                    messages=messages,
                    temperature=0.7,
                    max_tokens=1000
                )
                
                ai_response = response.choices[0].message.content
                ai_response = clean_response(ai_response)
                
                print("\r" + " "*20 + "\r", end="")
                print(f"AI: {ai_response}")
                
                messages.append({"role": "assistant", "content": ai_response})
                
            except Exception as e:
                print(f"\r[錯誤] {str(e)[:100]}")
                messages.pop()
                
        except KeyboardInterrupt:
            print("\n\n[中斷] 使用 exit 命令正常退出")
            continue
        except EOFError:
            print("\n再見！")
            break
    
    return "exit"

def main():
    while True:
        print("="*60)
        print("Llama AI 通用對話程式")
        print(f"時間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("="*60)
        
        # 選擇網路環境
        environment = choose_network_environment()
        
        # 尋找可用端點
        working_endpoint = find_working_endpoint(environment)
        
        if not working_endpoint:
            print("\n" + "="*60)
            print(f"錯誤：無法連接到任何 {NETWORK_CONFIG[environment]['name']} 端點")
            print("="*60)
            print("\n可能的原因：")
            if environment == "內網":
                print("1. 不在內網環境或 IP 設定錯誤")
                print("2. 內網 API 服務未啟動")
                print("3. 防火牆阻擋連接")
                print("\n建議：嘗試選擇外網環境")
            else:
                print("1. 外網 API 伺服器暫時離線")
                print("2. 網路連接問題")
                print("3. 防火牆或代理設定")
                print("\n建議：嘗試選擇內網環境或稍後再試")
            
            retry = input("\n是否重新選擇環境？(y/n，預設: y): ").strip().lower()
            if retry in ['n', 'no']:
                break
            continue
        
        print("\n" + "="*60)
        print(f"成功連接到: {working_endpoint['name']}")
        print("="*60)
        
        # 開始對話
        result = chat_session(working_endpoint, environment)
        
        if result == "exit":
            break
        elif result == "switch":
            continue  # 重新開始選擇環境

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n程式已退出")
    except Exception as e:
        print(f"\n[錯誤] {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)