# 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
from llama_parse import LlamaParse
from llama_index.core import SimpleDirectoryReader

load_dotenv(verbose=True)

chat_history = []
intreme = []

basic_format = """
「基本形」=学生や若手アナリスト向けの5つの項目を考えてみました：
1. 明確な投資仮説（Investment Thesis）
•	内容：なぜこの銘柄が過小評価されているのか、どのようにしてアルファを得るのかを、簡潔かつ明確に述べる。
•	理由：投資アイデアの核心を理解し、それを論理的に説明できているかを確認するため。

2. カタリストとタイミング
•	内容：株価の変動を促す可能性のあるイベント（決算、規制変更、新製品発売など）と、それがいつ起こるかの予測。
•	理由：アイデアの実現可能性と時間軸の現実性を評価するため。

3. バリュエーションとアップサイド／ダウンサイド
•	内容：DCFや相対バリュエーション、SOTPなどを用いた評価手法、目標株価、そしてリスク・リワードのバランス。
•	理由：アイデアの定量的な魅力や、リスクを踏まえた投資判断ができているかを確認するため。

4. 市場コンセンサスとの違い（Differentiated View）
•	内容：市場参加者の見解と自分の見解がどう異なるか、それを裏付けるデータや観点を示す。
•	理由：独自性のある分析ができているか、単なる後追いでないかを評価するため。

5. リスク評価と対応策
•	内容：主要リスク（例：業績未達、マージン圧力、マクロ要因など）と、それへの対応策や監視項目。
•	理由：リスクを適切に認識し、片側思考でなく投資判断をしているかを確認するため。
"""

# 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"])

def clean_text(text):
    # 改行とタブを削除
    return text.replace('\n', '').replace('\t', '').replace(' ', '')

def update_file_type(file):
    global filepath
    filepath = file
    if file is not None:
        # アップロードされたファイルの拡張子を取得
        _, ext = os.path.splitext(file.name)
        return ext
    return ""

# 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) > 4:
        #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):
    prompt = f"""以下の「レポート」を「基本形」に沿って正しく明記されているかをチェックし、漏れている項目を表示し、その内容をウェブから検索し、アドバイスしてください。最終的に内容がわかりやすかったか否かを教えてください。必ず、日本語で答えてください。
    「レポート」={msg}
    「基本形」={basic_format}
    """

    gresponse = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=[prompt]
    )
    summary = gresponse.text
    
    return summary

gresponse = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=["東証のtopixに関するニュースを教えてください。必ず、日本語で答えてください。"]
    )
tnews = gresponse.text

# Gradio UI with Blocks
with gr.Blocks(title="Fund Manager Academic 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 Academic BUDDY</span><img id="logo" src="https://www.ryhintl.com/images/ryhlogo/ryhlogo.png" width="64" height="64" alt="Logo"></div>')
    gr.Markdown("꧁ FM Acedemic Buddy ꧂ ")

    # 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="デフォルト")
                            file_input = gr.File(label="レポートをアップロード(PDF,pptx,xlsx,docx,txt,mdなど)")
                            file_type = gr.Textbox(label="ファイルタイプ")
                            dialog_output = gr.Chatbot(label="対話履歴", height=400, show_copy_all_button=True)
                            #dialog_output = gr.Textbox(label="対話履歴")
                            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):
    def generate_response(dtype, file_name,file_type, username):
        if (file_name == None):  # Should not happen if UI is correctly managed
            ret = []
            ret.append(["レポートをアップロードしてください。","Please Upload report."])
            yield ret
        else:
            intreme.append(["レポートを処理中...",tnews])
            yield intreme

        
            # LlamaParse Parser設定
            parser = LlamaParse(
                api_key=os.environ.get("LLAMA_API_KEY"),
                language="ja",
                user_prompt="提供された文書はスライドである。1ページにはファイルのタイトル、2ページは目次となっており、それ以降のページは各ページごとの副題を持つ。資料は、テキスト、図、表、グラフ、イラストを含む。",
                result_type="markdown"  # "markdown" または "text" が選択可能
            )
    
            file_extractor = {file_type: parser}
            documents = SimpleDirectoryReader(input_files=[file_name], file_extractor=file_extractor, recursive=False).load_data()
    
            texts = [doc.text for doc in documents]
        
            combined_text = clean_text(texts[0])
            msg = texts[0]

            bot_message = respond(file_type, msg, username)
            chat_history.append([msg, bot_message])
            intreme.clear()
            yield chat_history
    
    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

    file_input.change(fn=update_file_type, inputs=file_input, outputs=file_type)
        
    dialog_submit_btn.click(
        generate_response,
        #inputs=[dialog_ctype, dialog_query, dialog_output, current_username],
        inputs=[dialog_ctype, file_input, file_type, current_username],
        outputs=[dialog_output]
        #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)