import re

from openai import OpenAI

REQUIREMENT = """以下の箇条書きの文章群から、数値データや表データを含む箇所をそのまま抜き出してください。数値や表が含まれていない場合、その文章については None としてください。具体的には以下の条件に従ってください。

### 条件
-  数値データや表データを含む箇所をそのまま抜き出す
- 数値や表がない場合は None を返す
- 結果を箇条書きで表示する
- 複数データが含まれる場合はスラッシュ(/)区切りで一行に複数記す"""

EXAMPLE_QUESTION = """1. TSRの倒産集計における業種分類は、総務省の日本産業標準分類に準拠している。
2. 2024年1-8月の倒産は、ラーメン店が44件（前年同期比57.1％増）で、2009年からの統計では同期間で最多だった2020年の31件を大幅に上回っている。
3. 2009年以降の倒産動向をみると、ラーメン店は2013年に29件を記録した。その後、インバウンド（訪日外国人客）の来店増などで2016年は16件まで減少した"""

EXAMPLE_ANSWER = """1. None
2. 2024年1-8月の倒産は、ラーメン店が44件（前年同期比57.1％増）/ 2009年からの統計では同期間で最多だった2020年の31件
3. ラーメン店は2013年に29件を記録した/ 2016年は16件まで減少した"""


def get_data_info(client: OpenAI, facts, model_name='gpt-4o-mini-2024-07-18'):
    message = "\n".join([f'{i+1}. {q}' for i, q in enumerate(facts)])
    response = client.chat.completions.create(
        messages=[
            {"role": "user","content": REQUIREMENT},
            {"role": "user","content": EXAMPLE_QUESTION},
            {"role": "assistant","content": EXAMPLE_ANSWER},
            {"role": "user","content": message}
        ],
        model=model_name,
        temperature=0,
    )
    res_text = response.choices[0].message.content
    predicted_data = re.findall(r'\d{1,2}\. (.+)\n*', res_text)
    predicted_data = [[d_.strip() for d_ in d.split('/')] for d in predicted_data]
    predicted_data = '/'.join(['/'.join(d) for d in predicted_data if d != ['None']]).split('/')
    return predicted_data