clinical_flow_coordinator.py

from __future__ import annotations
from typing import Optional, Literal, Dict, Any, List
import torch, json, os, re
import asyncio
from concurrent.futures import ThreadPoolExecutor
from pydantic import BaseModel, Field, ValidationError
from transformers import AutoTokenizer, AutoModelForCausalLM

from domain_model import (
    ClinicalBlackboard, ClinicalState, WorkflowState, WorkflowPhase,
    ConsentStatus, PAStatus, HITLDecision, AgentDecision, WorkflowTask
)
from intent_engine import ClinicalIntentEngine, apply_emergency_override

intent_engine = ClinicalIntentEngine()

# Enable all 8 GPUs (0 through 7) for maximum parallel throughput.
if torch.cuda.is_available():
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"

_medgemma_client = None
_executor = ThreadPoolExecutor(max_workers=30)
_resolved_endpoint_cache: Dict[str, str] = {}
_dedicated_dns_cache: Dict[str, str] = {}

# =====================================================================
# 📌 1. Pydantic Structured Output Schemas (LLM Output Guardrails)
# =====================================================================
class TriageOutput(BaseModel):
    urgency: Literal["EMERGENT", "URGENT", "ROUTINE"] = Field(description="Clinical urgency level")
    reason: str = Field(description="Clinical reasoning for the urgency level")

class ChecklistOutput(BaseModel):
    missing_items: List[str] = Field(description="List of missing prep items (e.g., 'STAT Cr Lab', '18G IV')")
    reason: str = Field(description="Reason for the missing items")

class ConsentOutput(BaseModel):
    status: Literal["COMPLETED", "REQUIRED_MISSING"] = Field(description="Consent status. Must be exact.")
    reason: str = Field(description="Reasoning based on patient's mental capacity")

class PAOutput(BaseModel):
    status: Literal["APPROVED", "REQUIRED_MISSING"] = Field(description="Prior Authorization status.")
    reason: str = Field(description="Reasoning based on insurance and emergency intent.")

class ETAOutput(BaseModel):
    eta_minutes: int = Field(description="Estimated wait time in minutes")
    reason: str = Field(description="Reason for this time estimate")

class PriorityOutput(BaseModel):
    priority_score: float = Field(description="Dynamic priority score from 0.0 to 1.0")
    reason: str = Field(description="Reasoning for the priority score")

class CommOutput(BaseModel):
    messages: List[str] = Field(description="Specific dispatch messages for clinical roles in English")
    action_label: str = Field(description="A short 2-3 word label for the action button")

# =====================================================================
# 📌 2. Multi-GPU Model Pool (Native FP16) & Validation Logic
# =====================================================================
class MedGemmaClient:
    _instances = {}
    _locks = {}

    @classmethod
    def get_instance(cls, gpu_id: int, model_name="google/medgemma-27b-text-it"):
        mode = os.getenv("MODEL_BACKEND", "vertex").lower()
        if mode == "vertex":
            model_name = os.getenv("VERTEX_MEDGEMMA_MODEL_ID", "google_medgemma-1_5-4b-it-1771927636658")
        else:
            model_name = os.getenv("LOCAL_MODEL_NAME", model_name)
        key = f"{mode}:{gpu_id}:{model_name}"
        if key not in cls._instances:
            cls._instances[key] = cls(gpu_id, model_name=model_name, mode=mode)
            cls._locks[key] = asyncio.Lock()
        return cls._instances[key]

    def __init__(self, gpu_id: int, model_name="google/medgemma-27b-text-it", mode="vertex"):
        self.mode = mode
        self.model_name = model_name
        self.gpu_id = gpu_id
        self.device = f"cuda:{gpu_id}"
        self.vertex_client = None
        self.tokenizer = None
        self.model = None

        if self.mode == "vertex":
            self._init_vertex_client()
            return
        self._init_local_client()

    def _init_vertex_client(self):
        project_id = os.getenv("GOOGLE_CLOUD_PROJECT") or os.getenv("GCP_PROJECT_ID") or "project-1123e1e7-231b-45b8-981"
        location = os.getenv("VERTEX_MEDGEMMA_LOCATION") or os.getenv("VERTEX_LOCATION", "us-central1")
        if not project_id:
            raise RuntimeError("GOOGLE_CLOUD_PROJECT (or GCP_PROJECT_ID) is required for MODEL_BACKEND=vertex")
        try:
            from google import genai
            self.vertex_client = genai.Client(vertexai=True, project=project_id, location=location)
            print(f"☁️ Vertex AI client ready for model={self.model_name}, project={project_id}, location={location}")
        except Exception as e:
            raise RuntimeError(f"Failed to initialize Vertex AI client: {e}") from e

    def _extract_text_from_predict(self, payload: Dict[str, Any]) -> str:
        preds = payload.get("predictions", [])
        if not preds:
            return ""
        first = preds[0]
        if isinstance(first, str):
            return first
        for key in ["content", "text", "generated_text", "output", "response"]:
            val = first.get(key) if isinstance(first, dict) else None
            if isinstance(val, str):
                return val
        return json.dumps(first, ensure_ascii=False)

    def _structured_fallback(self, response_model: type[BaseModel]) -> Dict[str, Any]:
        fallback = {}
        for field_name, field_info in response_model.model_fields.items():
            ann = str(field_info.annotation)
            if "List" in ann:
                fallback[field_name] = []
            elif "float" in ann:
                fallback[field_name] = 0.5
            elif "int" in ann:
                fallback[field_name] = 0
            elif "Literal" in ann:
                fallback[field_name] = ann.split("'")[1]
            else:
                fallback[field_name] = "Fallback due to parsing error"
        return fallback

    def _parse_first_json_object(self, text: str) -> Dict[str, Any]:
        text = (text or "").strip()
        if not text:
            return {}
        # Prefer fenced JSON blocks if present.
        fence = re.search(r"```(?:json)?\s*(\{[\s\S]*?\})\s*```", text, re.IGNORECASE)
        candidate = fence.group(1) if fence else text

        # Find the first JSON object and decode only that object.
        start = candidate.find("{")
        if start < 0:
            return {}
        decoder = json.JSONDecoder()
        obj, _idx = decoder.raw_decode(candidate[start:])
        if isinstance(obj, dict):
            return obj
        return {}

    def _resolve_endpoint_resource(self, endpoint_hint: str, project_id: str, location: str) -> str:
        if endpoint_hint.startswith("projects/"):
            return endpoint_hint
        if endpoint_hint.isdigit():
            return f"projects/{project_id}/locations/{location}/endpoints/{endpoint_hint}"
        cache_key = f"{project_id}:{location}:{endpoint_hint}"
        if cache_key in _resolved_endpoint_cache:
            return _resolved_endpoint_cache[cache_key]

        from google.auth import default
        from google.auth.transport.requests import AuthorizedSession

        credentials, _ = default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        session = AuthorizedSession(credentials)
        list_url = f"https://aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/endpoints"
        filt = f'display_name="{endpoint_hint}"'
        resp = session.get(list_url, params={"filter": filt}, timeout=30)
        if resp.status_code < 400:
            endpoints = resp.json().get("endpoints", [])
            if endpoints:
                name = endpoints[0].get("name")
                if name:
                    _resolved_endpoint_cache[cache_key] = name
                    return name
        # Final fallback: treat hint as ID path fragment.
        return f"projects/{project_id}/locations/{location}/endpoints/{endpoint_hint}"

    def _get_dedicated_dns(self, endpoint_resource: str, location: str) -> str:
        cache_key = f"{location}:{endpoint_resource}"
        if cache_key in _dedicated_dns_cache:
            return _dedicated_dns_cache[cache_key]

        from google.auth import default
        from google.auth.transport.requests import AuthorizedSession

        credentials, _ = default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        session = AuthorizedSession(credentials)
        url = f"https://{location}-aiplatform.googleapis.com/v1/{endpoint_resource}"
        resp = session.get(url, timeout=30)
        if resp.status_code < 400:
            body = resp.json()
            dns = body.get("dedicatedEndpointDns") or body.get("dedicated_endpoint_dns") or ""
            if dns:
                dns = dns.replace("https://", "").strip("/")
                _dedicated_dns_cache[cache_key] = dns
                return dns
        return ""

    def _chat_vertex_endpoint_sync(self, system: str, user: str, response_model: type[BaseModel], max_tokens: int = 512) -> Dict[str, Any]:
        # Fallback path for one-click deployed endpoint when model API call is unavailable.
        endpoint = os.getenv("VERTEX_MEDGEMMA_ENDPOINT", "google_medgemma-1_5-4b-it-mg-one-click-deploy")
        project_id = os.getenv("GOOGLE_CLOUD_PROJECT") or os.getenv("GCP_PROJECT_ID") or "project-1123e1e7-231b-45b8-981"
        location = os.getenv("VERTEX_MEDGEMMA_LOCATION") or os.getenv("VERTEX_LOCATION", "us-central1")
        endpoint = self._resolve_endpoint_resource(endpoint, project_id, location)
        dedicated_dns = self._get_dedicated_dns(endpoint, location)

        from google.auth import default
        from google.auth.transport.requests import AuthorizedSession

        credentials, _ = default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
        session = AuthorizedSession(credentials)

        template_keys = {k: "..." for k in response_model.model_fields.keys()}
        schema_properties = response_model.model_json_schema().get("properties", {})
        prompt = (
            f"{system}\n\nCRITICAL: Return only one JSON object with exact keys {json.dumps(template_keys)}.\n"
            f"CRITICAL: Follow field types in this schema properties object: {json.dumps(schema_properties, ensure_ascii=False)}\n"
            f"CRITICAL: No markdown, no explanations, no trailing text.\n\n"
            f"User input:\n{user}"
        )
        urls: List[str] = []
        if dedicated_dns:
            urls.append(f"https://{dedicated_dns}/v1/{endpoint}:predict")
        urls.append(f"https://{location}-aiplatform.googleapis.com/v1/{endpoint}:predict")
        payloads = [
            {"instances": [{"prompt": prompt}], "parameters": {"temperature": 0.1, "maxOutputTokens": max_tokens}},
            {"instances": [{"content": prompt}], "parameters": {"temperature": 0.1, "maxOutputTokens": max_tokens}},
            {"instances": [{"inputs": prompt}], "parameters": {"temperature": 0.1, "maxOutputTokens": max_tokens}},
            {"instances": [prompt], "parameters": {"temperature": 0.1, "maxOutputTokens": max_tokens}},
        ]
        last_err = "unknown"
        out_text = ""
        for body in payloads:
            for url in urls:
                try:
                    response = session.post(url, json=body, timeout=60)
                except Exception as e:
                    last_err = str(e)
                    continue
                if response.status_code < 400:
                    out_text = self._extract_text_from_predict(response.json())
                    if out_text:
                        break
                else:
                    last_err = f"{response.status_code}: {response.text[:200]}"
            if out_text:
                break
        if not out_text:
            raise RuntimeError(f"Vertex endpoint predict failed for all payload variants on {endpoint}. last_error={last_err}")

        try:
            parsed = self._parse_first_json_object(out_text)
            validated = response_model(**parsed)
            return validated.model_dump()
        except Exception:
            # One retry with stricter repair prompt.
            repair_prompt = (
                "Convert the following content into a single valid JSON object only, strictly matching these keys: "
                f"{json.dumps(template_keys)}.\nContent:\n{out_text}"
            )
            repair_payloads = [
                {"instances": [{"prompt": repair_prompt}], "parameters": {"temperature": 0.0, "maxOutputTokens": max_tokens}},
                {"instances": [{"content": repair_prompt}], "parameters": {"temperature": 0.0, "maxOutputTokens": max_tokens}},
            ]
            for body in repair_payloads:
                for url in urls:
                    try:
                        response = session.post(url, json=body, timeout=60)
                    except Exception:
                        continue
                    if response.status_code >= 400:
                        continue
                    repaired_text = self._extract_text_from_predict(response.json())
                    try:
                        parsed = self._parse_first_json_object(repaired_text)
                        validated = response_model(**parsed)
                        return validated.model_dump()
                    except Exception:
                        continue
            return self._structured_fallback(response_model)

    def _init_local_client(self):
        print(f"🚀 Loading AI Brain on Dedicated GPU {self.gpu_id} (Native FP16)...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            device_map={"": self.device},
            torch_dtype=torch.float16
        )
        print(f"✅ GPU {self.gpu_id} Model Load Complete!")

    def _chat_local_sync(self, system: str, user: str, response_model: type[BaseModel], max_tokens: int = 512) -> Dict[str, Any]:
        # 🔥 BUG FIX: Extract only the keys to build a flat JSON template.
        # This prevents the LLM from outputting JSON schema syntax (like "properties" or "required").
        template_keys = {k: "..." for k in response_model.model_fields.keys()}
        template_str = json.dumps(template_keys, indent=2)
        
        system_with_schema = f"{system}\n\nCRITICAL INSTRUCTION: You MUST output ONLY a valid JSON object representing the data. DO NOT output a JSON schema definition. Your JSON must contain exactly these keys:\n{template_str}"
        
        prompt = f"<|im_start|>system\n{system_with_schema}\n<|im_end|>\n<|im_start|>user\n{user}\n<|im_end|>\n<|im_start|>assistant\n"
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        
        with torch.inference_mode():
            outputs = self.model.generate(
                **inputs, max_new_tokens=max_tokens, temperature=0.1, do_sample=True, pad_token_id=self.tokenizer.eos_token_id
            )
        response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).strip()
        
        response_clean = response.replace("<|im_end|>", "").strip()
        json_str = "{}"
        match = re.search(r'```(?:json)?\s*(\{.*\})\s*```', response_clean, re.DOTALL)
        if match: 
            json_str = match.group(1)
        else:
            match_simple = re.search(r'(\{.*\})', response_clean, re.DOTALL)
            if match_simple: 
                json_str = match_simple.group(1)
        
        try:
            parsed_dict = json.loads(json_str)
            # Pydantic Validation
            validated_data = response_model(**parsed_dict)
            return validated_data.model_dump()
        except (json.JSONDecodeError, ValidationError) as e:
            print(f"❌ [Validator Error] GPU {self.device} Output validation failed: {e}\nRaw Output: {json_str}")
            return self._structured_fallback(response_model)

    def _chat_vertex_sync(self, system: str, user: str, response_model: type[BaseModel], max_tokens: int = 512) -> Dict[str, Any]:
        # For MedGemma one-click deploy, endpoint predict is more reliable than model API.
        if os.getenv("VERTEX_MEDGEMMA_ENDPOINT"):
            return self._chat_vertex_endpoint_sync(system, user, response_model, max_tokens=max_tokens)
        try:
            from google.genai import types
        except Exception as e:
            raise RuntimeError(f"google-genai is required for Vertex backend: {e}") from e

        prompt = f"{system}\n\nUser input:\n{user}"
        schema = response_model.model_json_schema()
        try:
            response = self.vertex_client.models.generate_content(
                model=self.model_name,
                contents=prompt,
                config=types.GenerateContentConfig(
                    temperature=0.1,
                    max_output_tokens=max_tokens,
                    response_mime_type="application/json",
                    response_schema=schema,
                ),
            )
            text = getattr(response, "text", None) or "{}"
            parsed = json.loads(text)
            validated = response_model(**parsed)
            return validated.model_dump()
        except Exception as e:
            print(f"⚠️ Vertex model API failed. Falling back to endpoint predict: {e}")
            return self._chat_vertex_endpoint_sync(system, user, response_model, max_tokens=max_tokens)

    async def chat(self, system: str, user: str, response_model: type[BaseModel]) -> Dict[str, Any]:
        key = f"{self.mode}:{self.gpu_id}:{self.model_name}"
        async with self._locks[key]:
            loop = asyncio.get_running_loop()
            if self.mode == "vertex":
                return await loop.run_in_executor(_executor, self._chat_vertex_sync, system, user, response_model)
            return await loop.run_in_executor(_executor, self._chat_local_sync, system, user, response_model)

class MedGemmaAgent:
    def __init__(self, agent_id: str, agent_role: str, response_format_model: type[BaseModel], gpu_id: int):
        self.agent_id = agent_id
        self.agent_role = agent_role
        self.response_format_model = response_format_model
        self.client = MedGemmaClient.get_instance(gpu_id)

    async def run(self, bb: ClinicalBlackboard) -> AgentDecision:
        raise NotImplementedError("Subclasses must implement run()")

    async def _safe_chat(self, system: str, user: str) -> Dict[str, Any]:
        try:
            return await self.client.chat(system, user, self.response_format_model)
        except Exception as e:
            print(f"⚠️ {self.agent_id} model call failed: {e}")
            return {}

# =====================================================================
# 📌 3. Specialized Agents with Advanced Personas (Distributed across 8 GPUs)
# =====================================================================
class EmergencyPolicy:
    @staticmethod
    def adjust_blockers(state: ClinicalState):
        state.risk_notes = [b for b in state.readiness_blocking if b in ["NPO_UNKNOWN", "CONSENT_MISSING"]]
        state.readiness_blocking = [b for b in state.readiness_blocking if b not in ["NPO_UNKNOWN", "CONSENT_MISSING"]]


def _text_contains(text: str, keywords: List[str]) -> bool:
    t = (text or "").lower()
    return any(k in t for k in keywords)

class AcuityReasoner(MedGemmaAgent): 
    def __init__(self):
        super().__init__("triage_agent", "Triage Agent", TriageOutput, gpu_id=0) # GPU 0
    async def run(self, bb):
        cs = bb.clinical_state
        sys = """You are a Clinical Acuity Agent in a large tertiary hospital. 
Your role is to analyze patient vitals and procedure type to determine clinical urgency.
CRITICAL RULE: Output 'EMERGENT' if vitals are unstable (e.g., severe hypotension).
Example 1:
User: Vitals: BP 90/60, HR 120, severe pain. Procedure: Abdominal CT.
Assistant: {"urgency": "EMERGENT", "reason": "Patient shows signs of shock indicating an acute abdomen."}"""
        user = f"Vitals: {cs.vitals_summary}. Procedure: {cs.procedure_display}."
        parsed = await self._safe_chat(sys, user)
        vitals = (cs.vitals_summary or "").lower()
        proc = (cs.procedure_display or "").lower()
        if cs.intent == "emergent" or _text_contains(vitals + " " + proc, ["acute mental change", "stat", "shock", "resus"]):
            parsed = {"urgency": "EMERGENT", "reason": "STAT/emergency indicators require immediate prioritization."}
        elif _text_contains(vitals, ["fever", "pain", "confused", "tachy"]):
            parsed = {"urgency": "URGENT", "reason": "Symptoms indicate risk of near-term deterioration; same-day priority is needed."}
        else:
            parsed = {"urgency": "ROUTINE", "reason": "Vitals are stable and this case fits routine workflow."}
        cs.triage_level = parsed.get("urgency", "URGENT")
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="TRIAGE", decision=parsed)

class ReadinessReasoner(MedGemmaAgent):
    def __init__(self):
        super().__init__("checklist_agent", "Checklist Agent", ChecklistOutput, gpu_id=1) # GPU 1
    async def run(self, bb):
        cs = bb.clinical_state
        sys = """You are a Clinical Readiness Agent in a tertiary hospital. 
Verify prerequisites (NPO, Labs, IV) before a radiology exam.
CRITICAL RULES: If Procedure involves 'Contrast' and Labs show 'Cr' is missing/outdated, output 'STAT Cr Lab' in missing_items.
Example 1:
User: Intent: scheduled. Procedure: CT with Contrast. Labs: Cr 1.1 (40 days ago). Allergies: None.
Assistant: {"missing_items": ["STAT Cr Lab"], "reason": "Creatinine level is outdated."}"""
        user = f"Intent: {cs.intent}. Procedure: {cs.procedure_display}. Vitals: {cs.vitals_summary}. Labs: {cs.labs_summary}. Allergies: {cs.allergies}."
        parsed = await self._safe_chat(sys, user)
        procedure = (cs.procedure_display or "").lower()
        labs = (cs.labs_summary or "").lower()
        missing: List[str] = []
        if _text_contains(procedure, ["contrast", "apct"]) and _text_contains(labs, ["outdated", "unknown", "missing"]):
            missing.append("STAT Cr Lab")
        if cs.intent == "emergent" and _text_contains(procedure, ["contrast"]):
            # In emergent contrast CT, IV access is commonly a hard execution blocker.
            missing.append("18G IV")
        parsed = {
            "missing_items": missing,
            "reason": "Required readiness checks are complete." if not missing else "Critical prep for contrast/emergency workflow is missing."
        }
        cs.readiness_items_missing, cs.readiness_blocking = missing, missing
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="PREP", decision=parsed)

class ComplianceReasoner(MedGemmaAgent):
    def __init__(self):
        super().__init__("consent_agent", "Consent Agent", ConsentOutput, gpu_id=2) # GPU 2
    async def run(self, bb):
        cs = bb.clinical_state
        sys = """You are a Compliance & Legal Agent in a hospital. 
Determine if legal consent requirements are met based on the patient's mental capacity.
CRITICAL RULE: If Vitals mention 'Confused', 'altered', or 'unconscious', the patient lacks capacity. Output "REQUIRED_MISSING".
Example 1:
User: Procedure: MRI Brain. Vitals: Altered mental status, unresponsive. Comorbidities: None.
Assistant: {"status": "REQUIRED_MISSING", "reason": "Patient is unresponsive and lacks capacity. Surrogate consent is required."}"""
        user = f"Procedure: {cs.procedure_display}. Vitals: {cs.vitals_summary}. Comorbidities: {cs.comorbidities}."
        parsed = await self._safe_chat(sys, user)
        vitals = (cs.vitals_summary or "").lower()
        if _text_contains(vitals, ["confused", "altered", "unconscious", "acute mental change", "unresponsive"]):
            parsed = {"status": "REQUIRED_MISSING", "reason": "Patient lacks decisional capacity; surrogate consent is required."}
        else:
            parsed = {"status": "COMPLETED", "reason": "Patient has decision-making capacity; consent can be completed directly."}
        cs.consent_status = ConsentStatus.REQUIRED_MISSING if parsed.get("status") == "REQUIRED_MISSING" else ConsentStatus.COMPLETED
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="CONSENT", decision=parsed)

class FinancialClearanceReasoner(MedGemmaAgent):
    def __init__(self):
        super().__init__("pa_agent", "Prior Auth Agent", PAOutput, gpu_id=3) # GPU 3
    async def run(self, bb):
        cs = bb.clinical_state
        sys = """You are a Financial Clearance Agent in a hospital. 
Determine if Prior Authorization (PA) is required.
CRITICAL RULES:
1. Assume all demographics are perfectly recorded in the EHR. Never state they are missing.
2. If intent is 'emergent_stat', PA is bypassed. Output "APPROVED".
Example 1:
User: Patient: P9. Intent: scheduled. Insurance: HMO. Procedure: Knee MRI.
Assistant: {"status": "REQUIRED_MISSING", "reason": "MRI for HMO insurance typically requires prior authorization."}"""
        user = f"Patient: {cs.patient_id}. Intent: {cs.intent}. Insurance: {cs.insurance_plan}. Procedure: {cs.procedure_display}."
        parsed = await self._safe_chat(sys, user)
        ins = (cs.insurance_plan or "").lower()
        proc = (cs.procedure_display or "").lower()
        if cs.intent == "emergent":
            parsed = {"status": "APPROVED", "reason": "Emergency case; PA exception applies."}
        elif "hmo" in ins and _text_contains(proc, ["mri", "apct", "chest+apct", "bundle"]):
            parsed = {"status": "REQUIRED_MISSING", "reason": "HMO plan with higher-cost imaging requires prior authorization."}
        else:
            parsed = {"status": "APPROVED", "reason": "No additional prior authorization is required for this payer/order combination."}
        cs.pa_status = PAStatus.REQUIRED_MISSING if parsed.get("status") == "REQUIRED_MISSING" else PAStatus.APPROVED
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="PA", decision=parsed)

class ResourceCoordinationReasoner(MedGemmaAgent):
    def __init__(self):
        super().__init__("eta_agent", "ETA Agent", ETAOutput, gpu_id=4) # GPU 4
    async def run(self, bb):
        sys = """You are a Resource Dispatch Agent. Estimate wait time (ETA) in minutes.
Example 1:
User: Intent: scheduled. Procedure: Chest CT.
Assistant: {"eta_minutes": 15, "reason": "Scheduled procedures typically have minimal wait time upon arrival."}"""
        user = f"Intent: {bb.clinical_state.intent}. Procedure: {bb.clinical_state.procedure_display}."
        parsed = await self._safe_chat(sys, user)
        if bb.clinical_state.intent == "emergent":
            parsed = {"eta_minutes": 5, "reason": "STAT order requires immediate slot reallocation."}
        elif bb.clinical_state.intent == "same_day_addon":
            parsed = {"eta_minutes": 25, "reason": "Same-day add-on queue may introduce moderate delay."}
        else:
            parsed = {"eta_minutes": 15, "reason": "Scheduled workflow with typical wait-time range."}
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="ETA", decision=parsed)

class OperationalPrioritizationReasoner(MedGemmaAgent):
    def __init__(self):
        super().__init__("priority_agent", "Priority Agent", PriorityOutput, gpu_id=5) # GPU 5
    async def run(self, bb):
        sys = """You are a Dynamic Orchestrator Agent. Calculate priority score (0.0 to 1.0).
CRITICAL RULES: 'emergent_stat' = 0.9-1.0. 'same_day_addon' = 0.6-0.8. 'scheduled' = 0.1-0.5.
Example 1:
User: Intent: emergent_stat. Triage: EMERGENT.
Assistant: {"priority_score": 0.95, "reason": "STAT intent and EMERGENT triage dictate immediate priority."}"""
        user = f"Intent: {bb.clinical_state.intent}. Triage: {bb.clinical_state.triage_level}."
        parsed = await self._safe_chat(sys, user)
        intent = bb.clinical_state.intent
        triage = bb.clinical_state.triage_level or "URGENT"
        if intent == "emergent":
            score = 0.97
        elif intent == "same_day_addon":
            score = 0.72 if triage == "URGENT" else 0.64
        else:
            score = 0.35
        parsed = {"priority_score": score, "reason": f"Priority computed from Intent={intent} and Triage={triage}."}
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="PRIORITY", decision=parsed)

class CareTeamCommunicationAgent(MedGemmaAgent):
    def __init__(self):
        super().__init__("communication_agent", "Communication Agent", CommOutput, gpu_id=6) # GPU 6
    async def run(self, bb):
        cs = bb.clinical_state
        sys = """You are a Care Team Dispatcher. Translate blockers into actionable English messages.
CRITICAL RULES:
1. KEEP THE PATIENT ID EXACTLY AS IS (e.g., 'P3 - Kim, D.'). NEVER alter names.
2. If 'STAT Cr Lab' is blocked, ask the ward RN for immediate blood draw.
Example 1:
User: Exact Patient Name: P9 - Doe, J.\nLocation: ED\nProcedure: MRI\nBlockers: ['STAT Cr Lab']\nPA: APPROVED\nConsent: COMPLETED.
Assistant: {"messages": ["📣 [To Ward RN]: P9 - Doe, J. needs repeat creatinine blood draw before contrast use."], "action_label": "Lab Draw"}"""
        user = f"Exact Patient Name: {cs.patient_id}\nLocation: {cs.location}\nProcedure: {cs.procedure_display}\nBlockers: {cs.readiness_blocking}\nPA: {cs.pa_status.value}\nConsent: {cs.consent_status.value}."
        parsed = await self._safe_chat(sys, user)
        msgs: List[str] = []
        if cs.intent == "emergent":
            msgs.append(f"🚨 [To Radiology]: {cs.patient_id} requires immediate STAT slot allocation.")
        if "STAT Cr Lab" in cs.readiness_blocking:
            msgs.append(f"📣 [To Ward RN]: {cs.patient_id} needs repeat creatinine blood draw before contrast imaging.")
        if "18G IV" in cs.readiness_blocking:
            msgs.append(f"📣 [To ED RN]: Please establish an 18G IV line now for {cs.patient_id}.")
        if cs.consent_status.value == "required_missing":
            msgs.append(f"📣 [To Surrogate/Coordinator]: Surrogate consent workflow is needed for {cs.patient_id}.")
        if cs.pa_status.value == "required_missing":
            msgs.append(f"📣 [To Financial Clearance]: Prepare and submit PA packet for {cs.patient_id}.")
        if not msgs:
            msgs.append(f"✅ [To Transport Team]: {cs.patient_id} is ready for transport to imaging.")

        action_label = "Immediate Action" if cs.intent == "emergent" else ("PA Workflow" if cs.pa_status.value == "required_missing" else "Proceed to Transport")
        parsed = {"messages": msgs, "action_label": action_label}
        ui_msg = "\n\n".join(msgs) if isinstance(msgs, list) else str(msgs)
        return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="COMM", decision=parsed, ui_message=ui_msg)

class AdministrativeInterventionAgent:
    def __init__(self): self.agent_id, self.agent_role = "intervention_agent", "Administrative Intervention"
    async def run(self, bb): return AgentDecision(agent_id=self.agent_id, agent_role=self.agent_role, task_id="AUTO_RES")

class ClinicalReasoningCoordinator:
    def __init__(self, triage_agent, checklist_agent, consent_agent, pa_agent, eta_agent, priority_agent, comm_agent, intervention_agent):
        self.agents = {"triage_agent": triage_agent, "checklist_agent": checklist_agent, "consent_agent": consent_agent, "pa_agent": pa_agent, "eta_agent": eta_agent, "priority_agent": priority_agent, "communication_agent": comm_agent}

    async def step_stream(self, bb: ClinicalBlackboard):
        bb = intent_engine.evaluate_intent(bb)
        bb = apply_emergency_override(bb)
        if bb.clinical_state.intent == "emergent":
            EmergencyPolicy.adjust_blockers(bb.clinical_state)

        parallel_agents = ["triage_agent", "checklist_agent", "consent_agent", "pa_agent", "eta_agent"]
        
        async def run_agent(name):
            res = await self.agents[name].run(bb)
            return name, res

        tasks = [asyncio.create_task(run_agent(name)) for name in parallel_agents]
        
        for coro in asyncio.as_completed(tasks):
            name, result = await coro
            bb.agent_history.append(result)
            yield {"status": "agent_done", "agent": self.agents[name].agent_role, "decision": result.decision}

        for name in ["priority_agent", "communication_agent"]:
            result = await self.agents[name].run(bb)
            bb.agent_history.append(result)
            yield {"status": "agent_done", "agent": self.agents[name].agent_role, "decision": result.decision}

        cs = bb.clinical_state
        bb.workflow_state.phase = WorkflowPhase.BLOCKED if (cs.readiness_blocking or cs.pa_status.value == "REQUIRED_MISSING" or cs.consent_status.value == "REQUIRED_MISSING") else WorkflowPhase.EXECUTED
        yield {"status": "done", "bb": bb}