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

logger = logging.getLogger(__name__)

problem_prompt_template = PromptTemplate(
    input_variables=["responses", "internal_report"],
    template=(
        "You are a wellness analyst. You have the following user responses to health-related questions:\n"
        "{responses}\n\n"
        "You also have an internal analysis report:\n"
        "{internal_report}\n\n"
        "From these inputs, determine a 'problem severity percentage' for the user in the following areas: "
        "sleep, exercise, stress, and diet. "
        "Return your answer in JSON format with keys: sleep_problem, exercise_problem, stress_problem, diet_problem.\n"
        "Ensure severity percentages are numbers from 0 to 100.\n\n"
        "JSON Output:"
    )
)
problem_chain = LLMChain(llm=chat_model, prompt=problem_prompt_template)

def analyze_problems_with_chain(responses: Dict[str, str], internal_report: str) -> Dict[str, float]:
    responses_str = "\n".join(f"{q}: {a}" for q, a in responses.items())
    raw_text = problem_chain.run(responses=responses_str, internal_report=internal_report)
    try:
        start_idx = raw_text.find('{')
        end_idx = raw_text.rfind('}') + 1
        json_str = raw_text[start_idx:end_idx]
        problems = json.loads(json_str)
        for key in ["sleep_problem", "exercise_problem", "stress_problem", "diet_problem"]:
            problems.setdefault(key, 0.0)
        return {k: float(v) for k, v in problems.items()}
    except Exception as e:
        logger.error(f"Error parsing problem percentages from LLM: {e}")
        return {
            "sleep_problem": 0.0,
            "exercise_problem": 0.0,
            "stress_problem": 0.0,
            "diet_problem": 0.0
        }