# pip install gradio cohere python-dotenv numpy psycopg2-binary google-genai
import gradio as gr
import re
import hashlib
import sqlite3
import time

import cohere
import os
from dotenv import load_dotenv

import numpy as np
import psycopg2
from google import genai

import uuid
import base64
import requests
import asyncio
import random
import pandas as pd
from docx import Document

load_dotenv(verbose=True)

# Initialize Qdrant and Cohere clients
co = cohere.ClientV2(api_key=os.environ.get("COHERE_API_KEY"))
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

# Database Initialization and Utility Functions
def init_db():
    conn = sqlite3.connect('auth.db')
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            email TEXT UNIQUE NOT NULL,
            password TEXT NOT NULL,
            phone TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

init_db()

def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

def validate_username(username):
    if len(username) < 4:
        return "ユーザー名は4文字以上である必要があります"
    if not re.match("^[a-zA-Z0-9_]+$", username):
        return "ユーザー名には文字、数字、アンダースコアのみ使用できます"
    return None

def validate_email(email):
    if not re.match(r"^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+$", email):
        return "有効なメールアドレスを入力してください"
    return None

def validate_password(password):
    if len(password) < 8:
        return "パスワードは8文字以上でなければなりません"
    if not any(char.isdigit() for char in password):
        return "パスワードには少なくとも1つの数字を含める必要があります"
    if not any(char.isupper() for char in password):
        return "パスワードには少なくとも1つの大文字を含める必要があります"
    return None

def validate_phone(phone):
    if phone and not re.match(r"^\+?[0-9\s\-]+$", phone):
        return "有効な電話番号を入力してください"
    return None

def register_user(username, email, password, phone):
    conn = sqlite3.connect('auth.db')
    c = conn.cursor()
    c.execute("SELECT * FROM users WHERE username = ? OR email = ?", (username, email))
    if c.fetchone():
        conn.close()
        return False, "ユーザー名またはメールアドレスは既に存在します"

    hashed_pw = hash_password(password)
    c.execute(
        "INSERT INTO users (username, email, password, phone) VALUES (?, ?, ?, ?)",
        (username, email, hashed_pw, phone))
    conn.commit()
    conn.close()
    return True, "登録が成功しました"

def login_user(username, password):
    conn = sqlite3.connect('auth.db')
    c = conn.cursor()
    hashed_pw = hash_password(password)

    c.execute("SELECT * FROM users WHERE username = ? AND password = ?", (username, hashed_pw))
    user = c.fetchone()
    conn.close()

    if user:
        return True, user
    return False, "ユーザー名またはパスワードが無効です"

# Cohere and Gemini Helper Functions
def embed(input, input_type):
    response = co.embed(texts=input, model='embed-multilingual-v3.0', input_type=input_type, embedding_types=['ubinary'])
    return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary]

def summarize_text(long_text, username):
    if not long_text:
        return "要約するテキストがありません。"

    prompt = f"次の文章をuserとassistantに分けて的確に要約してください:\n{long_text}"
    gresponse = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=[prompt]
    )
    summary = gresponse.text

    #response = co.embed(
        #texts=[summary],
        #model="embed-multilingual-v3.0", 
        #input_type="search_document",
        #output_dimension=1024,
        #embedding_types=["float"],
    #)

    #embedding_str = ",".join(map(str, response.embeddings.float_[0]))

    #conn = psycopg2.connect(
        #dbname="smair",
        #user="smairuser",
        #password="smairuser",
        #host="www.ryhintl.com",
        #port=10629
    #)

    #cur = conn.cursor()
    #sql = "INSERT INTO dailog_logs (userid, content, embedding) VALUES (%s, %s, %s)"
    #cur.execute(sql, (username, summary, f"[{embedding_str}]"))
    #conn.commit()
    #cur.close()
    #conn.close()
    return summary


def generate_prompts(corp_name):
    # ここにプロンプト生成ロジックを実装
    # 例えば、以下のようにダミーのプロンプトを返す
    if len(corp_name) > 5:
        #Generate Prompts
        prompt = f"あなたは優秀なファンドマネージャーです。最近の{corp_name}の株の動きを調査して、アドバイスを求めるためのプロンプトを５つ程度、生成してください。必ず、生成されたプロンプトのみ出力してください。"
        gresponse = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[prompt]
        )
        prompts = gresponse.text
        return prompts
    else:
        return "企業名を入力してください。"

def respond(ctype, msg, username):
    conn = psycopg2.connect(
        dbname="smair",
        user="smairuser",
        password="smairuser",
        host="www.ryhintl.com",
        port=10629
    )
    cur = conn.cursor()

    query_embedding = embed([msg], 'search_query')[0].tolist()

    cur.execute(
        'SELECT content, 1 - (embedding <=> %s::vector) AS similarity FROM dailog_logs WHERE (1 - (embedding <=> %s::vector)) <> 0 ORDER BY similarity ASC',
        (query_embedding, query_embedding)
    )
    proof = [row[0] for row in cur.fetchall()]

    cur.execute(
        'SELECT content, 1 - (embedding <=> %s::vector) AS similarity FROM monitoring_dialog WHERE (1 - (embedding <=> %s::vector)) <> 0 ORDER BY similarity ASC',
        (query_embedding, query_embedding)
    )
    vector_resp = [row[0] for row in cur.fetchall()]

    if ctype == "デフォルト":
        message = f"{vector_resp}に基づいて{proof}を交えながら{msg}に対する答えを正確に出力してください。答えが見つからない場合は、ウェブから検索して答えてください。"
    else:
        message = f"{vector_resp}に基づいて{proof}を交えながら{msg}に対する答えを正確に関西弁で出力してください。答えが見つからない場合は、ウェブから検索して答えてください。"

    messages = [
        {"role": "system", "content": "あなたは、優秀なアシスタントです。"},
        {"role": "user", "content": message},
    ]

    response = co.chat(model="command-a-03-2025", messages=messages)
    bot_message = response.message.content[0].text
    
    cur.close()
    conn.close()
    return bot_message

# Gradio UI with Blocks
with gr.Blocks(title="Fund Manager Buddy", css="""footer {visibility: hidden;} #header {display: flex; justify-content: space-between; align-items: center; font-size: 24px; font-weight: bold;} #logo {width: 50px; height: 50px;} 
    .gradio-container {
        background-color: #f8f9fa;
    }
    .main {
        background-color: #f8f9fa;
    }
    .logo-container {
        position: absolute;
        top: 1px;
        left: 20px;
        z-index: 1000;
    }
    .logo-container img {
        height: 30px;
        width: auto;
    }
    .title {
        font-size: 1.5rem;
        font-weight: 700;
        color: #2c3e50;
        text-align: center;
        margin-bottom: 1.5rem;
    }
    .subtitle {
        font-size: 0.1rem;
        color: #7f8c8d;
        text-align: center;
        margin-bottom: 2rem;
    }
    .card {
        background: white;
        border-radius: 15px;
        padding: 2rem;
        box-shadow: 0 10px 20px rgba(0,0,0,0.1);
        margin-bottom: 1rem;
    }
    .success-message {
        color: #27ae60;
        text-align: center;
        margin-top: 1rem;
    }
    .error-message {
        color: #e74c3c;
        text-align: center;
        margin-top: 1rem;
    }
    .footer {
        text-align: center;
        margin-top: 1rem;
        color: #95a5a6;
        font-size: 0.8rem;
    }
    .avatar {
        width: 50px;
        height: 50px;
        border-radius: 50%;
        margin: 0 auto 1rem auto;
        display: block;
        object-fit: cover;
        border: 3px solid #4a90e2;
    }
    
    .avatar:hover {
        transform: scale(1.2);
        border: 3px solid #4a90e2;
    }
    .gr-button {
        width: 100%;
        border-radius: 10px;
        padding: 10px;
        background-color: #4a90e2;
        color: white;
        border: none;
        font-weight: 500;
        transition: all 0.3s;
    }
    .gr-button:hover {
        background-color: #357abd;
        color: lightyellow;
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
    }
    .gr-textinput, .gr-textbox {
        border-radius: 10px;
        padding: 10px;
        border: 1px solid #ced4da;
    }
""") as buddy:
    
    gr.HTML('<div id="header"><span>🛡️ FUND MANAGER BUDDY</span><img id="logo" src="https://www.ryhintl.com/images/ryhlogo/ryhlogo.png" width="64" height="64" alt="Logo"></div>')
    gr.Markdown("꧁ FM Buddy ꧂  ベクターDBに保存されている知識ベースのインベントリを使用してFMへの情報を共有します。")

    # State variables
    current_username = gr.State(None)
    current_user_info = gr.State(None)
    logged_in_state = gr.State(False)

    with gr.Column(elem_classes="main"):
        gr.HTML('<div><img src="https://www.ryhintl.com/images/fmgr.png" width="32" height="32" alt="Fund Manager Buddy"></div>')

        # Login/Register Section
        with gr.Column(visible=True, elem_id="login_register_section") as login_register_section:
            with gr.Tab("ログイン"):
                gr.HTML('<h5 class="title">お帰りなさい！</h5>')
                gr.HTML('<p style="text-align: center; margin-top: 1rem; font-size: 13px;">アカウントにアクセスするにはサインインしてください</p>')

                with gr.Column(elem_classes="card"):
                    login_username = gr.Textbox(label="ユーザー名", placeholder="ユーザー名を入力してください")
                    login_password = gr.Textbox(label="パスワード", type="password", placeholder="パスワードを入力してください")
                    login_status = gr.Markdown("")
                    login_btn = gr.Button("サインイン")

            with gr.Tab("アカウントを作成"):
                gr.HTML('<h1 class="title">アカウントを作成</h1>')
                gr.HTML('<p class="subtitle">今すぐ、参加して始めましょう！</p>')

                with gr.Column(elem_classes="card"):
                    reg_username = gr.Textbox(label="ユーザー名", placeholder="ユーザー名を選択してください")
                    reg_email = gr.Textbox(label="電子メール", placeholder="メールアドレスを入力してください")
                    reg_phone = gr.Textbox(label="電話番号 (オプション)", placeholder="+81-1234567890")
                    reg_password = gr.Textbox(label="パスワード", type="password", placeholder="パスワードを入力してください")
                    reg_confirm_password = gr.Textbox(label="パスワード再確認", type="password", placeholder="パスワードを再度入力してください")
                    register_status = gr.Markdown("")
                    register_btn = gr.Button("登録")

        # Dashboard Section
        with gr.Column(visible=False, elem_id="dashboard_section") as dashboard_section:
            welcome_message = gr.Markdown()
            gr.HTML('<p style="text-align: center; margin-top: 1rem; font-size: 10px;"">あなたは現在、アカウントにログインしています</p>')

            user_avatar = gr.HTML()

            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### QUERY EXAMPLES")
                    gen_query = gr.Textbox(label="企業名", placeholder="企業名を入力してください。", lines=1, info="例）トヨタ自動車")
                    generated_queries = gr.Textbox(label="クエリー例", lines=5)

                    gen_query.change(
                        generate_prompts,
                        inputs=[gen_query],
                        outputs=[generated_queries]
                    )


                with gr.Column(scale=3):
                    with gr.Tabs():
                        with gr.TabItem("ダイアログ"):
                            dialog_ctype = gr.Dropdown(["デフォルト", "関西弁"], label="会話形式", value="デフォルト")
                            dialog_query = gr.Textbox(label="クエリーを入力してください。", lines=5, value="トヨタ自動車の株を買いたいと思ってるんだけど、どう思いますか？ 尚、購入のための参考にすべき判断材料や基準を教えてください")
                            dialog_output = gr.Chatbot(label="対話履歴", height=400, show_copy_all_button=True)
                            dialog_submit_btn = gr.Button("生成")

                        with gr.TabItem("ダイアログ・サマリー"):
                            summary_input = gr.Textbox(label="要約したいユーザー・ダイアログを入力", lines=5)
                            summary_output = gr.Markdown(label="要約結果")
                            summary_submit_btn = gr.Button("要約")

                        with gr.TabItem("ファンド投資家向け最新情報"):
                            report_file = gr.Textbox(label="ファイル名")
                            report_input = gr.Textbox(label="投資家向けのアップデート情報を入力", lines=10, value='''
自分のプロフィル
山田 太郎
ファンドマネージャー
〇〇アセットマネジメント株式会社
連絡先: yamada@amgmt.com
                                                      
• パフォーマンス: 2025年第2四半期、当ファンドは手数料控除後で 3.50% のリターンを達成し、ベンチマークである MSCIワールド指数 を 1.20% 上回りました。
• 市場環境: この四半期は、技術革新の加速とインフレ圧力の緩和によって特徴づけられ、特に成長株に追い風となりました。
• 見通し: 引き続き慎重ながらも楽観的な見方をしており、構造的成長テーマと安定したキャッシュフローを持つ優良企業に注力しています。 

ベンチマークおよび同業他社とのリターン比較表
当ファンドは、情報技術セクターへの選好が奏功し、ベンチマークを上回るパフォーマンスを達成しました。同業他社と比較しても、相対的に堅調なリターンを確保しています。

パフォーマンス寄与度:
好調なパフォーマンスは主にテクノロジーセクターへのオーバーウェイト（比重超過）が牽引しました。特に、クラウドコンピューティングおよび人工知能関連企業への投資が大きく寄与しました。
一方で、公益事業セクターへのアンダーウェイト（比重不足）がわずかなマイナス要因となりましたが、全体的な影響は限定的でした。

資産配分を示す円グラフ:
    現在のポートフォリオは、株式が中心であり、特に情報技術、ヘルスケア、消費循環セクターに重点を置いています。
    地域別では、北米市場への配分が最も高く、次いで欧州、アジアとなっています。
リターンへの貢献度が高かった上位5銘柄と低かった下位5銘柄のリスト:
    （上位貢献銘柄の例：〇〇社、△△社、□□社など、具体的な企業名と貢献理由を記載）
    （下位貢献銘柄の例：××社、☆☆社など、具体的な企業名と影響理由を記載）
新たに組み入れた主要銘柄についての説明:
    今期は、再生可能エネルギー分野のリーディングカンパニーである「グリーンエネルギー・ソリューションズ」を新規に組み入れました。これは、長期的な脱炭素化トレンドと、同社の強固な技術力および市場シェアを評価したものです。
    現在の経済状況に関する簡単な分析:
    世界経済は緩やかな回復基調にあり、特にサービス業の活動が活発化しています。主要中央銀行はインフレ抑制と経済成長のバランスを取りながら金融政策を運営しており、金利の動向が引き続き注目されます。
    来四半期に向けたファンドのポジションニングについての説明:
    来四半期も引き続き、イノベーションを牽引する企業や、構造的な需要増加が見込まれるセクターに焦点を当てていきます。また、地政学的リスクやサプライチェーンの変動にも注意を払い、機動的なポートフォリオ調整を行う方針です。
                            
    運用資産総額 (AUM): 125.5 億ドル
    ファンド設定日: 2025年1月1日''')
                            report_output = gr.File(label="ダウンロード")
                            report_submit_btn = gr.Button("生成")

            logout_btn = gr.Button("ログアウト")

    gr.HTML('<p class="footer">© 2025 Fund Manager Buddy. All rights reserved.</p>')

    # Gradio Event Handlers

    # Login Logic
    def process_login(username, password, current_user_info_state):
        username_error = validate_username(username)
        password_error = validate_password(password)

        if username_error:
            return gr.update(value=f"<p class='error-message'>{username_error}</p>"), False, None, None, gr.update(visible=True), gr.update(visible=False)
        elif password_error:
            return gr.update(value=f"<p class='error-message'>{password_error}</p>"), False, None, None, gr.update(visible=True), gr.update(visible=False)
        else:
            success, result = login_user(username, password)
            if success:
                user_info = {
                    "id": result[0],
                    "username": result[1],
                    "email": result[2],
                    "phone": result[4]
                }
                return gr.update(value="<p class='success-message'></p>"), True, username, user_info, gr.update(visible=False), gr.update(visible=True)
            else:
                return gr.update(value=f"<p class='error-message'>{result}</p>"), False, None, None, gr.update(visible=True), gr.update(visible=False)

    login_btn.click(
        process_login,
        inputs=[login_username, login_password, current_user_info],
        outputs=[login_status, logged_in_state, current_username, current_user_info, login_register_section, dashboard_section]
    )

    # Register Logic
    def process_register(username, email, phone, password, confirm_password):
        errors = []
        username_error = validate_username(username)
        email_error = validate_email(email)
        password_error = validate_password(password)
        phone_error = validate_phone(phone)

        if username_error:
            errors.append(username_error)
        if email_error:
            errors.append(email_error)
        if password_error:
            errors.append(password_error)
        if phone_error:
            errors.append(phone_error)
        if password != confirm_password:
            errors.append("パスワードが一致しません。")

        if errors:
            error_html = "".join([f"<p class='error-message'>{e}</p>" for e in errors])
            return gr.update(value=error_html), gr.update(visible=True), gr.update(visible=False)
        else:
            success, message = register_user(username, email, password, phone)
            if success:
                return gr.update(value=f"<p class='success-message'>{message}</p>"), gr.update(visible=True), gr.update(visible=False)
            else:
                return gr.update(value=f"<p class='error-message'>{message}</p>"), gr.update(visible=True), gr.update(visible=False)

    register_btn.click(
        process_register,
        inputs=[reg_username, reg_email, reg_phone, reg_password, reg_confirm_password],
        outputs=[register_status, login_register_section, dashboard_section]
    )


    # Update Dashboard on Login Success
    def update_dashboard_ui(is_logged_in, username, user_info):
        if is_logged_in:
            welcome_msg = f'<h3 class="title">ようこそ！ {username} 様!</h3>'
            avatar_html = f'<img src="https://ui-avatars.com/api/?name=' + username + '&background=4a90e2&color=fff&size=200" class="avatar">'
            return welcome_msg, avatar_html, gr.update(visible=False), gr.update(visible=True)
        return "", "", gr.update(visible=True), gr.update(visible=False)

    logged_in_state.change(
        update_dashboard_ui,
        inputs=[logged_in_state, current_username, current_user_info],
        outputs=[welcome_message, user_avatar, login_register_section, dashboard_section]
    )

    # Logout Logic
    def process_logout():
        time.sleep(1)
        return False, None, None, gr.update(visible=True), gr.update(visible=False)

    logout_btn.click(
        process_logout,
        inputs=[],
        outputs=[logged_in_state, current_username, current_user_info, login_register_section, dashboard_section]
    )

    # Chat Response Logic
    def generate_response(ctype, msg, chat_history, username):
        if not username:  # Should not happen if UI is correctly managed
            return chat_history, gr.update(value="ログインしてください。")

        bot_message = respond(ctype, msg, username)
        chat_history.append([msg, bot_message])
        return chat_history, "" # Clear the input after sending
    
    def generate_report(report_file, report_input):
        if not report_input:
            return "アップデートがありません。"

        prompt = f"""以下のデータを元に[フォーマット]を修正してください。必ず、本文のみ出力してください。\n{report_input} 
[フォーマット]=
ファンド四半期投資家向け最新情報

投資家の皆様へ
件名: (ファンド名) 四半期投資家向け最新情報 - 期

1. エグゼクティブサマリー

    ・パフォーマンス:
    ・市場環境:
    ・見通し:

2. ファンドパフォーマンスレビュー
    ・(ベンチマークおよび同業他社とのリターン比較表)
    ・(パフォーマンス寄与度: 例: 「好調なパフォーマンスは主にテクノロジーセクターへのオーバーウェイト（比重超過）が牽引しました。一方、エネルギーセクターへのアンダーウェイト（比重不足）がわずかなマイナス要因となりました。」)

3. ポートフォリオレビュー
    ・(資産配分を示す円グラフ)
    ・(リターンへの貢献度が高かった上位5銘柄と低かった下位5銘柄のリスト)
    ・(新たに組み入れた主要銘柄についての説明)

4. 市場解説および見通し
    ・(現在の経済状況に関する簡単な分析)
    ・(来四半期に向けたファンドのポジションニングについての説明)

5. ファンド詳細
    ・運用資産総額 (AUM): XX.X 億ドル
    ・ファンド設定日: (日付)

6. 重要な開示事項と免責事項
    本資料は情報提供のみを目的としており、投資勧誘を目的とするものではありません。
    当ファンドの過去の運用実績は、将来の運用成果を保証するものではありません。
    投資には元本割れのリスクがあります。投資判断はご自身の責任において行ってください。
    本資料に記載されている情報は、作成時点のものであり、将来予告なく変更される場合があります。

敬具
(氏名 / ファンドマネージャー名) (役職) (ファンド運用会社名) (連絡先情報)
        """
        gresponse = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[prompt]
        )
        report_update = gresponse.text

        doc = Document()
        doc.add_paragraph(report_update)
        file_path = report_file

        target_file = ""

        if ".docx" in report_file:
            target_file = report_file
        else:
            target_file = report_file+".docx"

        doc.save(target_file)
        return target_file
        #return report_update
        
    dialog_submit_btn.click(
        generate_response,
        inputs=[dialog_ctype, dialog_query, dialog_output, current_username],
        outputs=[dialog_output, dialog_query]
    )

    report_submit_btn.click(
        generate_report,
        inputs=[report_file, report_input],
        outputs=[report_output]
    )

    # Summarize Logic
    def generate_summary(input_text, username):
        if not username:
            return "ログインしてください。"
        summary = summarize_text(input_text, username)
        return summary

    summary_submit_btn.click(
        generate_summary,
        inputs=[summary_input, current_username],
        outputs=summary_output
    )

buddy.launch(favicon_path="favicon.ico", show_api=False)