"Service to generate PPT structure."
import json
# from utils.logger import ServiceLogger
import logging

from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import MessagesPlaceholder
from langchain.tools import tool
from langchain_core.output_parsers import JsonOutputParser
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

from configs.config import OPENAI_API_KEY, OPENAI_MODEL_4O, OPENAI_MODEL_MINI
from services.ppt_generator.data_classes.presentation_outline import SlideLLM
from services.ppt_generator.data_classes.project import Project, Section, Slide

load_dotenv()


class ReviewerResponse(BaseModel):
    "this can be a list or a single item"
    section_title: str
    slide_title: str
    slide_suggestion: str


class PrivateEquityReviewer(BaseModel):
    persona: str = Field(..., title="This is a CEO or Private Equity Partner")
    industry: str
    sector: str
    pitch_type: str


@tool("Presentation_Reviewer")
def presentation_reviewer(
    persona_input: PrivateEquityReviewer, sections: list[Section]
):
    "Review the structure of the PPT"

    llm = ChatOpenAI(
        openai_api_key=OPENAI_API_KEY,
        temperature=1,
        model_name=OPENAI_MODEL_MINI,
    )

    chat_prompt = []
    if persona_input.persona == "CEO":
        chat_prompt.append(("system", "You are a CEO with 20 years of experience."))
    if persona_input.persona == "Private Equity Partner":
        chat_prompt.append(
            ("system", "You are a Private Equity Partner with 20 years of experience.")
        )
    else:
        chat_prompt.append(
            (
                "system",
                f"You are an expert in the {persona_input.industry} industry and the {persona_input.sector} sector.",
            )
        )

    chat_prompt.extend(
        [
            (
                "system",
                f"You specialize in the {persona_input.industry} industry and the {persona_input.sector} sector.",
            ),
            (
                "system",
                f"You are going to help review the structure of a {persona_input.pitch_type} presentation.",
            ),
            (
                "human",
                "The presentation has the following sections and slides: {input_sections}",
            ),
            (
                "human",
                """Give me feedback on how to improve the structure of the presentation at the slide level. \
                If you need to add a new slide in a section, then tile the slide title empty.""",
            ),
            ("system", "You will respond in a json object as {json_output}"),
        ]
    )

    prompt = ChatPromptTemplate.from_messages(chat_prompt)
    parser = JsonOutputParser(pydantic_object=ReviewerResponse)
    chain = prompt | llm | parser
    response = chain.invoke(
        {
            "input_sections": json.dumps(sections),
            "json_output": parser.get_format_instructions(),
        }
    )

    return json.dumps(response)


def adjust_slide(suggestion: ReviewerResponse) -> Section:
    "Adjust the slide based on the suggestion"
    # get the global project variable
    project = globals()["project"]
    # find the section from the suggestion
    section: Section = None
    for s in project.sections:
        if s.title == suggestion.section_title:
            section = s
            break
    # find the slide from the suggestion
    slide: Slide = None
    for s in section.slides:
        if s.title == suggestion.slide_title:
            slide = s
            break

    llm = ChatOpenAI(
        openai_api_key=OPENAI_API_KEY,
        temperature=0.5,
        model_name=OPENAI_MODEL_MINI,
    )

    chat_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                f"You are an investment banking expert in the {project.industry} industry and the {project.sector} sector.",
            ),
            (
                "system",
                """You are going to be provided with a suggestion on how to adjust a slide in a presentation and you need to \
            provide a new slide, with the suggestion included.""",
            ),
            ("human", f"The Current slide is: {slide.json()}"),
            ("human", f"The Suggestion is: {suggestion.slide_suggestion}"),
            ("system", "You will respond in a json object as {json_output}"),
        ]
    )

    parser = JsonOutputParser(pydantic_object=SlideLLM)

    chain = chat_prompt | llm | parser
    response = chain.invoke(
        {
            "json_output": parser.get_format_instructions(),
        }
    )

    # update the slide
    slide = response

    return response


model = ChatOpenAI(
    openai_api_key=OPENAI_API_KEY,
    temperature=0.1,
    model_name=OPENAI_MODEL_MINI,
)

tools = [presentation_reviewer]

agent = create_tool_calling_agent(
    llm=model,
    tools=tools,
    verbose=True,
)

agent_executor = AgentExecutor(
    agent=agent, tools=tools, verbose=True, max_iterations=15
)

project = Project.check_project_in_db(project_id="altlaw")