import json
import os

import requests
import streamlit as st
from anthropic import Anthropic
from anthropic.types import ToolUseBlock
from bs4 import BeautifulSoup
from googleapiclient.discovery import build
from rss_parser import RSSParser

if "rss" not in st.session_state:
    rss_url = "https://aws.amazon.com/about-aws/whats-new/recent/feed/"
    response = requests.get(rss_url)

    st.session_state["rss"] = RSSParser.parse(response.text)

rss = st.session_state["rss"]
items = rss.channel.items


def f1(**kwargs):
    st.session_state["selection"] = kwargs


def create_search_query(contents: str):

    client = Anthropic(
        api_key=os.environ.get("ANTHROPIC_API_KEY"),
    )

    SYSTEM_PROMPT = """あなたはGoogle検索を支援するアシスタントAIです。ユーザーから提供された資料に基づいて、関連情報を探すための検索クエリを生成する能力があります。"""

    USER_PROMPT_TEMPLATE = """
    <doc>
    {{ref-doc}}
    </doc>

    この記事を読んだ読者に、関連する情報のURLを共有してください。
    関連する情報はGoogle検索を利用して収集します。
    """

    # 2024/9/29修正
    # Toolを使用するように変更
    tool = {
        "name": "web_search",
        "description": "Google検索を行い情報を取得します。",
        "input_schema": {
            "type": "object",
            "properties": {
                "search_query": {
                    "type": "string",
                    "description": "検索クエリー",
                },
            },
            "required": ["search_query"],
        },
    }

    response = client.messages.create(
        max_tokens=1024,
        system=SYSTEM_PROMPT,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": USER_PROMPT_TEMPLATE.replace("{{ref-doc}}", contents),
                    }
                ],
            }
        ],
        temperature=0.5,
        model="claude-3-haiku-20240307",
        tools=[tool],
        tool_choice={"type": "tool", "name": "web_search"},
    )

    tool_use = list(filter(lambda x: x.type == "tool_use", response.content))[0]

    search_query = tool_use.input["search_query"]

    return search_query


def search(search_query: str):

    service = build("customsearch", "v1", developerKey=os.environ.get("GOOGLE_API_KEY"))
    cse = service.cse()
    res = cse.list(
        q=search_query,
        cx=os.environ.get("GOOGLE_CSE_ID"),
        # lr="lang_ja",
        num=3,
    ).execute()

    return res.get("items", [])


st.title("AWSのWhat's Newを解説くん")
st.write("←から選んでね")


with st.sidebar:

    for item in items:
        with st.container(border=True):
            st.write(item.title.content)

            # 2024/9/29修正
            # item.linkを参照していましたが、item.linksになったようです
            items = list(map(lambda x: x.content, item.links))

            kwargs = {
                "title": item.title.content,
                "pub_date": item.pub_date.content,
                "link": items[0],
            }

            st.button("おしえて", key=item.guid.content, on_click=f1, kwargs=kwargs)

# 5/26 https://tech.algomatic.jp/entry/2024/05/23/140219を参考にTypeScript型定義を使用

prompt = """
AWSのWhat's newを提示しますので、注意深く読み込んでください。

<AWSのWhat's new>
{{What's new}}
</AWSのWhat's new>

タスク
- 会話形式でわかりやすく解説してください。

登場人物
- user1: AWSスペシャリスト
- user2: 新入社員

ルール
- 回答は必ず日本語で行ってください。
- ハルシネーションは許されないので、What\'s newの内容からのみ回答を生成してください。
- また、回答付不要な出力は含めず、登場人物の会話のみを出力してください。
- 回答は以下のTypeScript型定義に従ってください。

```typescript
type Conversations = {
    conversations: Message[]
}

type Message = {
    actor: User,
    message: string
}

enum User {
    user1 = "user1",
    user2 = "user2"
}
```

例


{
    "conversations": [
        {
            "actor": "user1",
            "message": "AWSの新機能について説明しましょう。Mistral Smallというファウンデーションモデルが、Amazon Bedrock で利用可能になりました。"
        },
        {
            "actor": "user2",
            "message": "Mistral Smallとはどのようなモデルなのでしょうか?"
        },
        {
            "actor": "user1",
            "message": "Mistral Smallは、Mistral AIが提供する高性能な大規模言語モデルの1つです。Amazon Bedrock で利用できるようになりました。このモデルは、高速で低コストなパフォーマンスを提供し、大量の言語処理タスクに最適化されています。"
        },
        {
            "actor": "user2",
            "message": "Mistral Smallの主な特徴は何ですか?"
        },
        {
            "actor": "user1",
            "message": "Mistral Smallの主な特徴は以下の通りです:情報保持能力に優れた retrieval-augmented generation (RAG) 機能を搭載コーディングタスクに強い性能を発揮英語、フランス語、ドイツ語、スペイン語、イタリア語など、多言語に対応安全性の高い機能も備えている"
        },
        {
            "actor": "user2",
            "message": "Mistral Smallはどのようなユースケースに適しているのでしょうか?"
        },
        {
            "actor": "user1",
            "message": "Mistral Smallは、大量の言語処理タスクを効率的に処理できるため、分類、カスタマーサポート、テキスト生成などの用途に適しています。また、コードの生成、レビュー、コメント付与などのソフトウェア開発タスクにも活用できます。"
        },
        {
            "actor": "user2",
            "message": "Mistral Smallはどこで利用できるのですか?"
        },
        {
            "actor": "user1",
            "message": "Mistral Smallは、Amazon Bedrock で利用できるようになりました。Amazon Bedrock コンソールからアクセスして、Mistral Smallを使ったアプリケーションの構築を始められます。"
        }
    ]
}



会話形式でわかりやすく解説してください。
"""


if "selection" in st.session_state:
    kwargs = st.session_state["selection"]

    response = requests.get(kwargs["link"])
    soup = BeautifulSoup(response.text, features="html.parser")

    # 2024/5/21修正
    # What's Newの構成が変わり、mainをチェックするように
    contents = soup.main.get_text()

    st.subheader(kwargs["title"])
    st.write(kwargs["pub_date"])

    with st.spinner("ちょっとまってね"):

        client = Anthropic(
            api_key=os.environ.get("ANTHROPIC_API_KEY"),
        )

        # 2024/9/29修正
        # Toolを使用するように変更
        tool = {
            "name": "conversation_memory",
            "description": "会話履歴をメモリーに保存します",
            "input_schema": {
                "type": "object",
                "properties": {
                    "conversations": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "actor": {
                                    "type": "string",
                                    "description": "会話の発言者。user1またはuser2",
                                },
                                "answer": {
                                    "type": "string",
                                    "message": "会話の内容",
                                },
                            },
                        },
                    },
                },
                "required": ["conversations"],
            },
        }

        message = client.messages.create(
            max_tokens=1024,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt.replace("{{What's new}}", contents),
                        }
                    ],
                }
            ],
            temperature=0.5,
            model="claude-3-haiku-20240307",
            tools=[tool],
            tool_choice={"type": "tool", "name": "conversation_memory"},
        )

    try:
        tool_use = message.content[0]
        assert isinstance(tool_use, ToolUseBlock)
        conversations = tool_use.input

        for c in conversations["conversations"]:
            if c["actor"] == "user1":
                with st.chat_message("ai"):
                    st.write(c["message"])

            if c["actor"] == "user2":
                with st.chat_message("human"):
                    st.write(c["message"])

    except Exception as e:
        st.write("ごめんなさい。エラーです。")
        with st.expander("Haikuの回答："):
            st.write(message.content)

    st.link_button("ブログを確認", kwargs["link"])

    st.divider()

    search_query = create_search_query(contents=contents)
    # st.write(search_query)

    if search_query:
        search_result = search(search_query=search_query)
        # st.write(search_result)

        st.markdown("おすすめ記事")
        for item in search_result:
            with st.container(border=True):
                st.markdown(f'##### {item["title"]}')
                st.markdown(item["htmlSnippet"], unsafe_allow_html=True)
                st.link_button(item["formattedUrl"], item["link"])