# chain_summary.py
import json
from typing import Dict
from langchain import PromptTemplate, LLMChain
from models import chat_model

final_prompt_template = PromptTemplate(
    input_variables=["report", "problems", "recommendation"],
    template=(
        "Based on the following information:\n"
        "Report:\n{report}\n\n"
        "Problem Severity Percentages:\n{problems}\n\n"
        "Recommended Packages:\n{recommendation}\n\n"
        "Generate a short summary suitable for video narration that synthesizes this information."
    )
)
final_chain = LLMChain(llm=chat_model, prompt=final_prompt_template)

def generate_final_summary(report: str, problems: Dict[str, float], recommendation: str) -> str:
    summary = final_chain.run(
        report=report,
        problems=json.dumps(problems),
        recommendation=recommendation
    )
    return summary.strip()

shorten_prompt_template = PromptTemplate(
    input_variables=["final_summary"],
    template=(
        "Shorten the following summary to make it concise and engaging for video narration. "
        "Ensure all key points remain intact:\n\n"
        "{final_summary}\n\n"
        "Shortened Summary:"
    )
)
shorten_chain = LLMChain(llm=chat_model, prompt=shorten_prompt_template)

def shorten_summary(final_summary: str) -> str:
    shortened = shorten_chain.run(final_summary=final_summary)
    return shortened.strip()