#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Llama API 外網連接程式
使用外網端點進行 AI 對話
"""

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

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

# 外網 API 端點配置
ENDPOINTS = [
    {
        "name": "Llama 通用端點",
        "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"]
    }
]

# 外網 API 端點配置（僅包含實際存在的端點）

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=10):
    """測試端點是否可用"""
    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(" ✓ 可用")
        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(f" ✗ 錯誤")
        return False

def find_working_endpoint():
    """尋找可用的端點"""
    print("\n正在測試可用端點...")
    print("-" * 50)
    
    for endpoint in ENDPOINTS:
        if test_endpoint(endpoint):
            return endpoint
    
    return None

def chat_session(endpoint_info):
    """對話主程式"""
    print("\n" + "="*60)
    print("Llama AI 對話系統")
    print("="*60)
    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("-"*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() == '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

def main():
    print("="*60)
    print("Llama AI 外網對話程式")
    print(f"時間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("="*60)
    
    # 尋找可用端點
    working_endpoint = find_working_endpoint()
    
    if not working_endpoint:
        print("\n" + "="*60)
        print("錯誤：無法連接到任何 API 端點")
        print("="*60)
        print("\n可能的原因：")
        print("1. API 伺服器暫時離線")
        print("2. 網路連接問題")
        print("3. 防火牆或代理設定")
        print("\n建議：")
        print("1. 稍後再試（10-30分鐘後）")
        print("2. 檢查網路連接")
        print("3. 確認可以訪問 https://llama.theaken.com")
        sys.exit(1)
    
    print("\n" + "="*60)
    print(f"成功連接到: {working_endpoint['name']}")
    print("="*60)
    
    # 開始對話
    chat_session(working_endpoint)

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)