import gradio as gr
import sys
import cohere
import os
import json
import requests
from dotenv import load_dotenv

# --- 初期設定 ---

# 環境変数のロード
load_dotenv(verbose=True)
COHERE_API_KEY = os.environ.get("COHERE_API_KEY")

# Cohereクライアントの初期化とエラーチェック
if COHERE_API_KEY:
    try:
        co = cohere.ClientV2(api_key=COHERE_API_KEY)
    except Exception as e:
        print(f"Cohereクライアント初期化エラー: {e}")
        co = None
else:
    print("⚠️ COHERE_API_KEYが設定されていません。AI機能は利用できません。")
    co = None

# --- 設定: 許可されたIPアドレスのリストを外部から取得する関数 ---
# SQLコマンドのエンコードは維持
API_URL = "https://www.ryhintl.com/dbjson/getjson?sqlcmd=select * from allowed_ips"

def fetch_allowed_ips():
    """指定されたAPIエンドポイントから許可IPリストを動的に取得する。"""
    print(f"許可IPリストを外部URLから取得中: {API_URL}")
    
    try:
        response = requests.get(API_URL, timeout=10) # タイムアウトを設定
        response.raise_for_status() # HTTPエラー（4xx, 5xx）をチェック
        
        data = response.json()
        print('data:', data) # 取得したデータを確認
        
        # 新しいJSON構造 [{'allowed': 'IP1'}, {'allowed': 'IP2'}, ...] に合わせてパース
        allowed_ips = []
        if isinstance(data, list):
            for item in data:
                if isinstance(item, dict) and "allowed" in item:
                    ip = item["allowed"].strip()
                    if ip:
                        allowed_ips.append(ip)
        
        if not allowed_ips:
            print("警告: 外部DBから有効なIPアドレスが取得されませんでした。")
        
        print("allowed:",allowed_ips)
        return allowed_ips

    except requests.exceptions.RequestException as e:
        print(f"APIリクエストエラー: {e}")
        return []
    except json.JSONDecodeError as e:
        print(f"JSONパースエラー: {e}")
        return []
    except Exception as e:
        print(f"予期せぬエラー: {e}")
        return []

# スクリプト実行時に許可IPリストを初期化
ALLOWED_IPS = fetch_allowed_ips()
print('ALLOWED_IPS:',ALLOWED_IPS)

# --- Cohere 処理関数 ---
def req_process(prompt: str):
    if not co:
        return "エラー: COHERE_API_KEYが設定されていないため、AI機能は実行できません。"
        
    try:
        res = co.chat(
            model="command-a-03-2025",
            messages=[
                {
                    "role": "user",
                    "content": f'{prompt}に答えてください',
                }
            ],
        )
        # Cohere V2 SDKのレスポンス構造に合わせる
        final = res.message.content[0].text
        return final
    except Exception as e:
        return f"Cohere APIエラーが発生しました: {e}"


# --- IPアドレス取得・チェック関数 ---
# 関数定義から不要な 'ins' 引数を削除（修正済）
def get_client_ip(request: gr.Request) -> tuple[str, gr.update, gr.update]:
    """
    gr.RequestオブジェクトからクライアントのIPアドレスを抽出し、許可リストと比較します。
    戻り値: (ip_outputへのメッセージ, outs_outputの表示状態を制御するgr.update, main_content_groupの表示状態を制御するgr.update)
    """
    if not request:
        return "IPアドレス情報が利用できませんでした。", gr.update(visible=False), gr.update(visible=False)

    # 1. 実際のクライアントIPを特定する
    client_ip = request.client.host
    forwarded_for = request.headers.get("x-forwarded-for")
    
    # 外部アクセスの場合にチェックすべきIP (X-Forwarded-Forの最初のIP)
    check_ip = client_ip
    if forwarded_for:
        # X-Forwarded-Forヘッダーはコンマ区切りで複数のIPを含む場合があるため、最初のIPを使用
        check_ip = forwarded_for.split(',')[0].strip()

    # 2. IPアドレスのチェック
    # ALLOWED_IPSはグローバル変数として既に定義されている
    if check_ip or client_ip in ALLOWED_IPS:
        # --- IPが許可されている場合 ---
        print('checkips:',check_ip,ALLOWED_IPS)
        
        if forwarded_for:
            display_msg = f"✅ アクセス許可: IPアドレスが許可リストに含まれています。\nクライアントIP (X-Forwarded-For): {check_ip}\n直接ホストIP: {client_ip}"
        else:
            display_msg = f"✅ アクセス許可: IPアドレスが許可リストに含まれています。\nクライアントIP (直接接続): {client_ip}"
        
        # outs_output（表示）、main_content_group（表示）
        return display_msg, gr.update(visible=True), gr.update(visible=True)
    else:
        # --- IPが拒否された場合 ---
        
        error_msg = f"❌ アクセス拒否: あなたのIPアドレス ({check_ip}) は許可されていません。\nアクセスを許可するには、外部DBの設定を確認してください。"
        
        # outs_output（非表示）、main_content_group（非表示）
        return error_msg, gr.update(visible=False), gr.update(visible=False)


# --- Gradio UIの構築 ---
with gr.Blocks(title="アクセスIPチェッカー（IP制限）") as demo:
    gr.Markdown(
        """
        # Gradio アクセスIPアドレス確認ツール（IP制限付き）
        
        このアプリケーションは、外部DBから取得したIPアドレスリストに登録されたIPからのみ、
        「制限されたメインコンテンツ」セクション全体を表示します。
        
        ---
        """
    )

    # 1. IPチェック結果を表示するテキストボックス (非許可の場合もエラーメッセージ表示のために常に存在)
    ip_output = gr.Textbox(
        label="アクセスステータス (IPチェック結果)",
        interactive=False,
        lines=4,
        value="ページロード時にアクセスチェックを実行します..."
    )

    # 2. 制限をかけたい主要なコンテンツを gr.Group で囲む
    with gr.Group(visible=False) as main_content_group:
        gr.Markdown(
            """
            ## 制限されたメインコンテンツエリア (Cohere Chat機能)
            このセクションは、許可されたIPアドレスからのみ表示されています。
            """
        )

        outs_output = gr.Textbox(
            label="接続ステータス",
            interactive=False,
            lines=1,
            value="接続ステータスがアクティブです。",
        )

        ins_output = gr.Textbox(
            label="プロンプト入力",
            interactive=True,
            lines=4
        )

        req_output = gr.Textbox(
            label="リクエスト結果",
            interactive=False,
            lines=8
        )

        exec_button = gr.Button("🚀 実行", variant="primary")

        exec_button.click(
            fn=req_process,
            inputs=[ins_output], 
            outputs=[req_output] 
        )
    
    # ページのロードが完了したときに、自動的にIPアドレスのチェック関数を呼び出す
    demo.load(
        fn=get_client_ip,
        inputs=[], # 入力コンポーネントは不要
        outputs=[ip_output, outs_output, main_content_group] 
    )

# デモの起動 (ローカルで実行する場合)
if __name__ == "__main__":
    print(f"許可されたIPアドレス (外部DBからの取得結果): {ALLOWED_IPS}")
    print("Gradioアプリケーションを起動します。")
        
    demo.launch()