scenario_builder.py

import json
import hashlib
import re
from typing import Dict, Any, List

class ScenarioNode:
    def __init__(self, question: str, options: Dict[str, Any], correct_answer: str = None, feedback: str = None):
        self.question = question
        self.options = options
        self.correct_answer = correct_answer
        self.feedback = feedback

    def to_dict(self):
        return {
            "question": self.question,
            "options": self.options,
            "correct_answer": self.correct_answer,
            "feedback": self.feedback
        }

class ScenarioBuilder:
    def __init__(self, ai_engine):
        self.ai_engine = ai_engine
        self.scenarios = {}

    async def create_scenario_from_query(self, query: str, num_questions: int = 5) -> Dict[str, Any]:
        prompt = f"""
You are a medical training assistant.

Based on the topic: "{query}", generate a {num_questions}-question interactive medical scenario. Each question must have:

- a `question` (string)
- an `options` object with keys "A"–"D"
- a `correct_answer` ("A", "B", "C", or "D")
- a `feedback` string

Respond ONLY with this exact JSON format:

[
  {{
    "question": "...",
    "options": {{
      "A": "...",
      "B": "...",
      "C": "...",
      "D": "..."
    }},
    "correct_answer": "A|B|C|D",
    "feedback": "..."
  }},
  ...
]

⚠️ No markdown. No prose. No headings. Just valid JSON array. Your job is to act like a JSON API.
"""

        try:
            # Use Groq for accurate structure
            raw_response = self.ai_engine._generate_with_groq(
                query=query,
                refined_prompt=prompt
            )

            # Clean any wrapping or garbage
            start = raw_response.find("[")
            end = raw_response.rfind("]") + 1
            json_block = raw_response[start:end].strip()

            print("🧪 Raw Groq output preview:", json_block[:200])

            scenario_data = json.loads(json_block)

            # Turn into ScenarioNode objects
            nodes = []
            for q in scenario_data:
                node = ScenarioNode(
                    question=q["question"],
                    options=q["options"],
                    correct_answer=q["correct_answer"],
                    feedback=q["feedback"]
                )
                nodes.append(node)

            scenario_id = hashlib.md5(query.encode()).hexdigest()
            self.scenarios[scenario_id] = nodes

            return {
                "scenario_id": scenario_id,
                "questions": [n.to_dict() for n in nodes]
            }

        except json.JSONDecodeError as e:
            print("❌ JSON parse error:", e)
            print("⚠️ Raw Groq output:", raw_response[:500])
            return {
                "error": "Failed to parse AI-generated scenario. Invalid JSON format.",
                "details": str(e)
            }

        except Exception as e:
            print("❌ Scenario generation error:", e)
            return {
                "error": "Failed to generate scenario.",
                "details": str(e)
            }

    def get_scenario(self, scenario_id: str) -> List[ScenarioNode]:
        return self.scenarios.get(scenario_id)