orchestrator.py

"""
============================================================
Mentallico Orchestrator - The Brain
============================================================
هنا بنحط كل المكونات مع بعض:

flow على رسالة جديدة:
  1. (لو صوت) → STT
  2. Classifier ensemble → prediction
  3. Conversation manager → should_diagnose?
  4a. لو لسه → probing/chat response
  4b. لو نعم → final diagnosis + recommendations
  5. حفظ الـ history وارجاع response

الـ orchestrator هو الـ entry point الوحيد للـ API.
"""
from typing import Optional, Dict, List
from dataclasses import dataclass

from config import CLASSIFICATION_MODELS, LABELS, MODELS_DIR, DIAGNOSIS_CONFIG
from classifier_models import EnsembleDiagnoser
from conversation_manager import ConversationManager
from rag_engine import get_rag
from speech_to_text import get_transcriber
import os
import torch


class MentallicoBrain:
    """المنسّق الرئيسي للنظام كله"""

    def __init__(self, primary_model_path: str = None, secondary_model_path: str = None,
                 enable_stt: bool = True):
        # Default paths
        if primary_model_path is None:
            primary_model_path = os.path.join(MODELS_DIR, "primary_model.pt")
        if secondary_model_path is None:
            secondary_model_path = os.path.join(MODELS_DIR, "secondary_model.pt")

        # ===== Classifier Ensemble =====
        self.diagnoser = EnsembleDiagnoser(CLASSIFICATION_MODELS, LABELS)
        primary = primary_model_path if os.path.exists(primary_model_path) else None
        secondary = secondary_model_path if os.path.exists(secondary_model_path) else None
        if primary or secondary:
            self.diagnoser.load_pretrained(primary_path=primary, secondary_path=secondary)
        else:
            print("⚠ No trained models found. Only emotion model + RAG will work.")
            self.diagnoser.load_pretrained()  # just emotion

        # ===== RAG =====
        try:
            self.rag = get_rag()
        except Exception as e:
            print(f"⚠ RAG init failed: {e}")
            self.rag = None

        # ===== Conversation Manager =====
        self.conv_mgr = ConversationManager()

        # ===== STT =====
        self.transcriber = get_transcriber() if enable_stt else None

    # ============= Sessions =============
    def create_session(self) -> str:
        return self.conv_mgr.create_session()

    def reset_session(self, session_id: str):
        self.conv_mgr.reset(session_id)

    # ============= Audio =============
    def process_audio(self, audio_path: str, language: str = None) -> dict:
        """يحوّل صوت لـ text"""
        if self.transcriber is None:
            return {"error": "STT not enabled"}
        return self.transcriber.transcribe(audio_path, language)

    # ============= Main Flow =============
    def process_message(self, session_id: str, user_text: str,
                        user_audio_path: str = None) -> dict:
        """
        أهم function - يدير التفاعل كله.

        ترجع:
          {
            "type": "chat" | "probing" | "diagnosis",
            "response": "...",
            "diagnosis": "..." (لو type==diagnosis),
            "confidence": 0.85,
            "urgency": "normal" | "critical",
            "session_id": "...",
            "internal": {...} (للـ debugging)
          }
        """
        # ===== 1. Audio processing لو موجود =====
        audio_info = None
        if user_audio_path:
            audio_result = self.process_audio(user_audio_path)
            if audio_result.get('text'):
                user_text = audio_result['text']
                audio_info = audio_result
            elif not user_text:
                return {
                    "type": "error",
                    "response": "Sorry, I couldn't understand the audio. Could you type your message?",
                    "error": audio_result.get('error', 'Audio processing failed')
                }

        if not user_text or not user_text.strip():
            return {"type": "error", "response": "Empty input received."}

        # ===== 2. Classifier prediction (على النص الجديد + التراكمي) =====
        # نحلل النص التراكمي عشان نحصل على signal أقوى
        combined_text = self.conv_mgr.get_combined_text(session_id) + " " + user_text
        prediction = self.diagnoser.predict(combined_text.strip())

        # نحفظ الرسالة مع التنبؤ
        self.conv_mgr.add_user_message(session_id, user_text, prediction=prediction)

        # ===== 3. هل نشخّص؟ =====
        decision = self.conv_mgr.should_diagnose(session_id, prediction)

        # ===== 4. ولّد الرد المناسب =====
        history = [
            {"role": m.role, "content": m.content}
            for m in self.conv_mgr.get_session(session_id).messages
        ]

        if decision['diagnose']:
            # ===== التشخيص النهائي =====
            label = decision['label']
            confidence = decision['confidence']
            urgency = decision.get('urgency', 'normal')

            if self.rag is not None:
                response_text = self.rag.final_diagnosis_response(
                    diagnosis=label,
                    confidence=confidence,
                    history=history,
                    urgency=urgency
                )
            else:
                response_text = self._fallback_diagnosis_text(label, confidence)

            self.conv_mgr.add_assistant_message(session_id, response_text)
            self.conv_mgr.confirm_diagnosis(session_id, label, confidence)

            return {
                "type": "diagnosis",
                "response": response_text,
                "diagnosis": label,
                "confidence": round(confidence, 3),
                "urgency": urgency,
                "session_id": session_id,
                "audio_info": audio_info,
                "internal": {
                    "decision_reason": decision['reason'],
                    "consistency": decision.get('consistency'),
                    "current_prediction": prediction,
                }
            }
        else:
            # ===== لسه بنجمّع معلومات =====
            # لو الموديل عنده توقع قوي نعمل probing، غير كده chat عادي
            confidence = prediction['confidence']
            if confidence >= 0.5 and self.rag is not None:
                response_text = self.rag.probing_response(
                    user_input=user_text,
                    history=history,
                    predicted_label=prediction['label'],
                    confidence=confidence
                )
                response_type = "probing"
            else:
                if self.rag is not None:
                    response_text = self.rag.chat_response(user_text, history)
                else:
                    response_text = self._fallback_chat_text()
                response_type = "chat"

            self.conv_mgr.add_assistant_message(session_id, response_text)

            return {
                "type": response_type,
                "response": response_text,
                "diagnosis": None,
                "confidence": round(confidence, 3),
                "urgency": "normal",
                "session_id": session_id,
                "audio_info": audio_info,
                "internal": {
                    "decision_reason": decision['reason'],
                    "progress": decision.get('progress'),
                    "current_prediction": prediction,
                }
            }

    def _fallback_diagnosis_text(self, label: str, confidence: float) -> str:
        return (f"Based on our conversation, the analysis suggests **{label}** "
                f"(confidence: {confidence:.0%}).\n\n"
                f"⚠ The full RAG system is not available, so detailed recommendations "
                f"can't be generated. Please consult a mental health professional for "
                f"a proper evaluation and treatment plan.")

    def _fallback_chat_text(self) -> str:
        return ("I'm here to listen. Could you tell me more about how you've been feeling lately? "
                "When did these feelings start, and how have they been affecting your daily life?")

    # ============= Knowledge base management =============
    def add_pdf(self, pdf_path: str) -> dict:
        if self.rag is None:
            return {"status": "error", "message": "RAG not initialized"}
        return self.rag.add_pdf_to_knowledge_base(pdf_path)


# ===== Global singleton (lazy) =====
_brain = None


def get_brain() -> MentallicoBrain:
    global _brain
    if _brain is None:
        _brain = MentallicoBrain()
    return _brain