import json

import pytest
from langchain import PromptTemplate, LLMChain
from langchain.chat_models import ChatOpenAI
from FOMCTexts.prompts import PROMPT_EXTRACT_DATE


PROMPT_DATE = PromptTemplate.from_template(PROMPT_EXTRACT_DATE)
date_extractor = LLMChain(prompt=PROMPT_DATE, llm=ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo'))


@pytest.mark.parametrize(
    "statement, expected_dict_date",
    [
        ('Summarize in two paragraphs the monetary policy outlook discussed by the participants during the meeting of June 2023', {"year": "2023", "month": "06"}),
        ("What were the key takeaways from the conference in September 2022?", {"year": "2022", "month": "09"}),
        ("Provide an update on the project status in March 2023.", {"year": "2023", "month": "03"}),
        ("Who attended the meeting on May 2021?", {"year": "2021", "month": "05"}),
        ("Please summarize the findings of the study conducted in January 2022.", {"year": "2022", "month": "01"}),
        ("Who were the participants in the event on July 2020?", {"year": "2020", "month": "07"}),
        ("Provide an overview of the report in August 2019.", {"year": "2019", "month": "08"}),
        ("Who presented the findings of the research in November 2021?", {"year": "2021", "month": "11"}),
        ("What were the discussions during the meeting in December 2024?", {"year": "2024", "month": "12"})
    ]
)
def test_date_extraction(statement, expected_dict_date):
    dict_date = json.loads(date_extractor.run(statement).replace('\n', '').replace(' ', ''))
    for k in dict_date:
        assert dict_date[k] == expected_dict_date[k]