import os
import json
from pathlib import Path
from typing import Annotated
from autogen import AssistantAgent, UserProxyAgent
from autogen.coding import LocalCommandLineCodeExecutor
import gradio as gr
from autogen import ConversableAgent
from autogen import register_function
import mysql.connector
import random
import requests
from groq import Groq
from dotenv import load_dotenv
import urllib.parse


#tool_resp = ""
tool_topic = ""

js = """
function createGradioAnimation() {
    var container = document.createElement('div');
    container.id = 'gradio-animation';
    container.style.fontSize = '2em';
    container.style.fontWeight = 'bold';
    container.style.textAlign = 'center';
    container.style.marginBottom = '20px';

    var text = 'wiki分析';
    for (var i = 0; i < text.length; i++) {
        (function(i){
            setTimeout(function(){
                var letter = document.createElement('span');
                var randomColor = "#" + Math.floor(Math.random() * 16777215).toString(16);
                letter.style.color = randomColor;
                letter.style.opacity = '0';
                letter.style.transition = 'opacity 0.5s';
                letter.innerText = text[i];

                container.appendChild(letter);

                setTimeout(function() {
                    letter.style.opacity = '1';
                }, 50);

                // Blink the text 3 times
                for (var j = 0; j < 3; j++) {
                    setTimeout(function() {
                        letter.style.opacity = '0';
                    }, 500 + j * 1000);
                    setTimeout(function() {
                        letter.style.opacity = '1';
                    }, 1000 + j * 1000);
                }
            }, i * 250);
        })(i);
    }

    var gradioContainer = document.querySelector('.gradio-container');
    gradioContainer.insertBefore(container, gradioContainer.firstChild);

    return 'Animation created';
}
"""

load_dotenv(verbose=True)

conn = mysql.connector.connect(
        host=os.environ.get("HOST"),
        user=os.environ.get("USER_NAME"),
        password=os.environ.get("PASSWORD"),
        port=os.environ.get("PORT"),
        database=os.environ.get("DB"),
        ssl_disabled=True,
        connection_timeout=60,
        use_pure=True
    )

cursor = conn.cursor(dictionary=True)

def get_rounrobin():
    select_one_data_query = "select api from agentic_apis_count order by counts ASC"
    cursor.execute(select_one_data_query)
    result = cursor.fetchall()
    first_api = result[0]['api']
    return first_api
    
# MySQLに接続
def get_api_keys():
    token = get_rounrobin()
    os.environ["GROQ_API_KEY"] = token
    
    return token

# Configure Groq
config_list = [{
    "model": "llama-3.3-70b-versatile",
    "api_key": os.environ["GROQ_API_KEY"],
    "api_type": "groq"
}]


# Create a directory to store code files from code executor
work_dir = Path("coding")
work_dir.mkdir(exist_ok=True)
code_executor = LocalCommandLineCodeExecutor(work_dir=work_dir)

# Define wiki tool
def get_wiki_content(topic):
    """Get the info for some material"""
    
    global tool_topic
    print("parms:",tool_topic)
    tmp_url = urllib.parse.unquote('https://ja.wikipedia.org/api/rest_v1/page/summary/'+tool_topic)
    #print("tmp_url:",tmp_url)
    data = requests.get(tmp_url)
    #print('request:',data.content.decode("utf-8"))
    # 元のデータ
    datas = json.loads(data.content.decode("utf-8"))
    #print('jload:',datas)


    # 指定された形式に変換
    #revenue_data = {item["title"]: {"revenue": item["extract"]} for item in datas}
    #print("revenue data:",revenue_data)
    
    jdump = json.dumps({
        "title": datas["title"],
        "content": datas["extract"],
    })
    #print("jdump:",jdump)
    return jdump

# Create an AI assistant that uses the kpi tool
assistant = AssistantAgent(
#assistant = ConversableAgent(
    name="groq_assistant",
    system_message="""あなたは、次のことができる役に立つAIアシスタントです。
    - 情報検索ツールを使用する
    - 結果を分析して自然言語のみで説明する""",
    llm_config={"config_list": config_list}
)

# Create a user proxy agent that only handles code execution
user_proxy = UserProxyAgent(
#user_proxy = ConversableAgent(
    name="user_proxy",
    human_input_mode="NEVER",
    code_execution_config={"work_dir":"coding", "use_docker":False},
    max_consecutive_auto_reply=2,
    #llm_config={"config_list": config_list}
)


# Register weather tool with the assistant
@user_proxy.register_for_execution()
@assistant.register_for_llm(description="extractの内容")
def revenue_analysis(
    title: Annotated[str, "title"],
    topic: Annotated[str, "topic"]
) -> str:
    #global tool_topic
    print("parameters:",title,topic)
    wiki_details = get_wiki_content(title)
    wikis = json.loads(wiki_details)
    print("myRevenue:",wikis)

    return f"{wikis['title']}の内容は{wikis['content']}"


def get_wiki_and_respond(prompt,topic):
    get_api_keys()

    print("mytopic-is:",topic)
    global tool_topic
    tool_topic = topic

    # Start the conversation
    resp = user_proxy.initiate_chat(
        assistant,
        message=f"""3つのことをやってみましょう:
        1. {prompt}の内容をtoolを利用して抽出します。
        2. toolを利用して抽出された内容を詳しく分析します。
        3. 日本語で説明してください。
        """
    )
    
    total_tokens = resp.cost['usage_including_cached_inference']['llama-3.3-70b-versatile']['total_tokens']
    
    groq_assistant_contents = [entry['content'] for entry in resp.chat_history if entry['role'] == 'user' and entry['name'] == 'groq_assistant']

    usages = "使用トークン数: "+str(total_tokens)
    return groq_assistant_contents,usages

# Create Gradio interface
iface = gr.Interface(
    js=js,
    fn=get_wiki_and_respond,
    inputs=[gr.Textbox(label="プロンプト",value="アユタヤ王朝とは何ですか？"),gr.Textbox(label="topic",value="アユタヤ王朝")],
    outputs=[gr.Textbox(label="結果"),gr.Textbox(label="Usageデータ")],
    title="資料の分析",
    description="プロンプトを入力してデータを取得し、内容を分析します。",
    submit_btn="実行",
    clear_btn="クリア",
    flagging_mode="never"
)

iface.launch()