# -*- coding: utf-8 -*-

import os
import sys
from transformers import AutoModelForCausalLM, AutoTokenizer
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate

# 環境変数からHugging FaceのトークンとGroq APIキーを取得
os.environ['HF_TOKEN'] = os.getenv('HUGGING_FACE_TOKEN')
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

# Hugging Faceのモデルとトークナイザーをロード
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.getenv('HF_TOKEN'))
model = AutoModelForCausalLM.from_pretrained(model_id, token=os.getenv('HF_TOKEN'))

# Groqを使ったLangChainのChatGroqインスタンスを初期化
chat = ChatGroq(model_name=model_id, tokenizer=tokenizer)

# 検索キーワード
keyword = sys.argv[1]

# output.txtからテキストを読み込む
with open("output1.txt", "r", encoding="utf-8") as f:
    output1_text = f.read()

# プロンプトを設定してGroq APIを使用して見出しを生成
system_prompt = "あなたはライティングの専門家です。PREP法に則って見出しを構成してください。"
user_prompt = f'（検索キーワード）の感情分析を行った結果に基づいて、興味や関心を引きつける「（検索キーワード）を含めたタイトル」、興味や関心を引きつける「見出し」を生成してください。そしてそれぞれの見出しに沿って「小見出し」を適宜生成してください。最後に、記事を締めくくる「<h2>まとめ</h2>」のセクションを追加してください。タイトル、見出し、小見出しはPREP法に基づいて一貫性のあるものとし、記事の前半部分で（検索キーワード）の検索意図を満たす記事構成内容にしてください。（テキスト）で使用されているキーワードが（検索キーワード）と関連していると判断できる場合、それらのキーワードを、見出しおよび小見出しに使用してください。ただし、見出しおよび小見出しに固有名詞のキーワードは使用しないでください。タイトル、見出しおよび小見出しは順に「<h1>（検索キーワード）を含めたタイトル</h1>」「<h2>1.見出し</h2>」「<h3>1-1 小見出し とは</h3>、<h3>1-2 小見出し とは</h3>」「<h2>2.見出し</h2>」「<h3>2-1 小見出し とは</h3>、<h3>2-2 小見出し とは</h3>」というようにいくつかの見出しと小見出しで構成してください。ただし、見出しは最大4つ、小見出しは最大8つまでの構成とします。各見出し、小見出しに説明文は不要です。タイトル、見出し、そして小見出しは検索者が求める情報を整理して記載するようにしてください。トピッククラスターモデルを意識し、（テキスト）内の異なるトピックに明確な関連性がある場合は異なるトピックをまとめ「記事を一つ作成=タイトルを一つ作成」します。異なるトピックに明確な関連性がない場合は「記事を適宜複数作成=タイトルを適宜複数作成」してください。（テキスト）および（検索キーワード）から検索ニーズを考慮してください。なお、似たような見出しは避けてください。テキスト：\"{output1_text}\" 検索キーワード：\"{keyword}\"'
prompt = ChatPromptTemplate.from_messages([("system", system_prompt), ("human", user_prompt)])
result = chat.invoke(prompt)

# output2.txtにまとめたテキストを書き込む
with open("output2.txt", "w", encoding="utf-8") as f:
    f.write(result.content)

# output2.txtの存在確認と内容表示
if os.path.exists("output2.txt"):
    print("output2.txt exists.")
    with open("output2.txt", "r", encoding="utf-8") as f:
        content = f.read()
        print("Content of output2.txt:")
        print(content)
else:
    print("output2.txt does not exist.")