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

recommend_prompt_template = PromptTemplate(
    input_variables=["problems"],
    template=(
        "Given the following problem severity percentages:\n"
        "{problems}\n\n"
        "Using these rules:\n"
        "- If sleep_problem > 70: Recommend Sleep Improvement Package\n"
        "- If stress_problem > 70: Recommend Stress Reduction Package\n"
        "- If exercise_problem > 70: Recommend Exercise Enhancement Package\n"
        "- If all problems are between 30 and 70: Recommend Balanced Wellness Package\n"
        "- If no severe problems: Recommend General Wellness Package\n\n"
        "What are the recommended wellness packages?"
    )
)
recommend_chain = LLMChain(llm=chat_model, prompt=recommend_prompt_template)

def generate_recommendations(problems: Dict[str, float]) -> str:
    recommendations = recommend_chain.run(problems=json.dumps(problems))
    return recommendations.strip()