## This module will be used to pull the non summarized PR data from the database
import logging

from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_openai import ChatOpenAI
from pydantic import BaseModel

from configs.config import OPENAI_API_KEY, OPENAI_MODEL_35
from services.PressRelease.PRSources.data_models.PR_Model import PRModel

logging.basicConfig(level=logging.INFO)


## LLM DataClasses
class RelevantPR(BaseModel):
    relevant_PR: bool


class PRSummary(BaseModel):
    summary: str


class RelevantInfo(BaseModel):
    relevant_info: dict[str, str]


class PRSummarizer:
    llm = ChatOpenAI(
        temperature=0.1, model_name=OPENAI_MODEL_35, api_key=OPENAI_API_KEY
    )

    def __init__(self):
        pass

    async def summarize(self, pr: PRModel) -> PRModel:
        """Process PRs"""
        logging.info(f"Summarizing PR: {pr}")
        # is PR relevant? #TODO: Implement a proprietary Model for this
        pr.relevant_PR = await self._determine_relevance(pr)

        if pr.relevant_PR:
            pr.summary = await self._summarize_text(pr)
            pr.relevant_info = await self._get_relevant_info(pr)

            if pr.relevant_info.get("Company Name"):
                pr.company_name = pr.relevant_info.get("Company Name")
            elif pr.relevant_info.get("Company"):
                pr.company_name = pr.relevant_info.get("Company")

            if pr.relevant_info.get("Website"):
                pr.company_url = pr.relevant_info.get("Website")

        pr.processed = True
        pr.update_modified()
        return pr

    async def _determine_relevance(self, pr: PRModel) -> bool:
        """Determine if the PR is relevant"""

        chat_temp = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """You are an expert analyzer of press releases as an expert Investment Banker. 
             You have 20 years of experience. You are going to analyze the headline and tell us if it is a relevant press release""",
                ),
                ("human", pr.title),
                (
                    "system",
                    "Is this a relevant press release? Return True or False as a JSON output.",
                ),
                ("system", "The JSON output is structured like this {json_output}"),
            ]
        )
        parser = JsonOutputParser(pydantic_object=RelevantPR)

        chain = chat_temp | self.llm | parser
        try:
            content = await chain.ainvoke(
                {"json_output": parser.get_format_instructions()}
            )
            return content["relevant_PR"]
        except Exception as e:
            logging.error(f"Exception while determining relevance: {e}")

        return False

    async def _summarize_text(self, pr: PRModel) -> str:
        """Summarize the PR text"""
        chat_temp = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """You are an expert analyzer of press releases as an expert Investment Banker. 
             You have 20 years of experience. You are going to summarize the press release text and return relevant information that a i-banker would care about""",
                ),
                ("human", pr.text),
                ("system", "The JSON output is structured like this {json_output}"),
            ]
        )
        parser = JsonOutputParser(pydantic_object=PRSummary)

        chain = chat_temp | self.llm | parser
        try:
            content = await chain.ainvoke(
                {"json_output": parser.get_format_instructions()}
            )
            return content["summary"]
        except Exception as e:
            logging.error(f"Exception while summarizing text {e}")

        return ""

    async def _get_relevant_info(self, pr: PRModel) -> dict:
        """Get relevant information from the PR"""
        chat_temp = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    """You are an expert analyzer of press releases as an expert Investment Banker. 
             You have 20 years of experience. You are going to extract relevant information from the press release text""",
                ),
                ("human", pr.text),
                (
                    "system",
                    """You are going to provide key value pairs of relevant information from the press release text.\
             We need to get information like company name, acquire name, partners, EBITDA, revenue, price, or similar information.""",
                ),
                ("system", "The JSON output is structured like this {json_output}"),
            ]
        )
        parser = JsonOutputParser(pydantic_object=RelevantInfo)

        chain = chat_temp | self.llm | parser
        try:
            content = await chain.ainvoke(
                {"json_output": parser.get_format_instructions()}
            )
            return content["relevant_info"]
        except Exception as e:
            logging.error(f"Exception thrown while getting relevant info: {e}")

        return {}