client.py

import base64
import io
import json
import os
import re

import modal
import numpy as np
import scipy.io.wavfile as wavfile
import soundfile as sf
import torch
from gtts import gTTS
from kokoro import KPipeline
from PIL import Image
from transformers import AutoTokenizer, VitsModel

# Import local handlers to enable seamless fallback when Modal is unavailable or offline mode is checked
import local_mode


def is_local_mode() -> bool:
    """Dynamically checks if Local Mode is enabled via environment variables."""
    return os.environ.get("LOCAL_MODE", "false").lower() == "true"


def transcribe_audio(audio_bytes: bytes) -> str:
    """Transcribe symptoms in Hindi. Calls local Whisper if local mode is active or Modal fails."""
    if is_local_mode():
        print("[Client] Running local transcription...")
        return local_mode.local_transcribe(audio_bytes)

    try:
        print("[Client] Calling Modal run_transcription...")
        # Get deployed Modal app function
        f = modal.Function.lookup("village-health-screener", "run_transcription")
        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
        return f.remote(audio_b64)
    except Exception as e:
        print(
            f"[Client WARNING] Modal transcription call failed: {e}. Falling back to local mode."
        )
        return local_mode.local_transcribe(audio_bytes)


def describe_image(pil_image: Image.Image) -> str:
    """Describe symptom visible photo. Calls local SmolVLM if local mode is active or Modal fails."""
    if is_local_mode():
        print("[Client] Running local vision analysis...")
        prompt = (
            "Describe any visible medical symptoms, skin conditions, wounds, rashes, swelling, "
            "or abnormalities in this image. Be specific and clinical. "
            "If there are no visible symptoms, say so clearly."
        )
        return local_mode.local_vision(pil_image, prompt)

    try:
        print("[Client] Calling Modal run_vision...")
        f = modal.Function.lookup("village-health-screener", "run_vision")

        # Convert PIL to bytes
        img_byte_arr = io.BytesIO()
        pil_image.save(img_byte_arr, format="JPEG")
        img_bytes = img_byte_arr.getvalue()
        image_b64 = base64.b64encode(img_bytes).decode("utf-8")

        prompt = (
            "Describe any visible medical symptoms, skin conditions, wounds, rashes, swelling, "
            "or abnormalities in this image. Be specific and clinical. "
            "If there are no visible symptoms, say so clearly."
        )
        return f.remote(image_b64, prompt)
    except Exception as e:
        print(
            f"[Client WARNING] Modal vision call failed: {e}. Falling back to local mode."
        )
        prompt = (
            "Describe any visible medical symptoms, skin conditions, wounds, rashes, swelling, "
            "or abnormalities in this image. Be specific and clinical. "
            "If there are no visible symptoms, say so clearly."
        )
        return local_mode.local_vision(pil_image, prompt)


def run_triage(symptoms_text: str, visual_description: str = None) -> dict:
    """Produce clinical triage report. Calls local Nemotron-Mini-4B if local mode is active or Modal fails."""
    system_prompt = (
        "You are a clinical triage assistant for ASHA health workers in rural India. "
        "Your job is to analyze patient symptoms and visible signs and produce a structured triage report. "
        "Always respond in this exact JSON format with these six fields:\n"
        "{\n"
        '  "likely_condition": "string",\n'
        '  "severity": integer from 1 to 5 where 1 is mild and 5 is emergency,\n'
        '  "immediate_actions": ["list of strings"],\n'
        '  "refer_to_doctor": boolean,\n'
        '  "refer_reason": "string, empty if false",\n'
        '  "followup_days": integer\n'
        "}\n"
        "Be conservative. When in doubt, recommend referral. "
        "Keep language simple enough for a health worker with basic training."
    )

    user_message = f"Spoken Hindi Symptoms: {symptoms_text}\n"
    if visual_description:
        user_message += f"Visual Symptom Observation: {visual_description}\n"

    raw_response = ""
    if is_local_mode():
        print("[Client] Running local Nemotron inference...")
        raw_response = local_mode.local_inference(
            system_prompt, user_message, max_tokens=600
        )
    else:
        try:
            print("[Client] Calling Modal run_inference...")
            f = modal.Function.lookup("village-health-screener", "run_inference")
            raw_response = f.remote(system_prompt, user_message, max_tokens=600)
        except Exception as e:
            print(
                f"[Client WARNING] Modal inference call failed: {e}. Falling back to local mode."
            )
            raw_response = local_mode.local_inference(
                system_prompt, user_message, max_tokens=600
            )

    # Parse and validate the response
    return _parse_and_validate_json(raw_response)


def _parse_and_validate_json(raw_text: str) -> dict:
    """Clean, parse, and validate JSON output defensively using regex and structured fallbacks."""
    fallback_dict = {
        "likely_condition": "Requires Medical Evaluation",
        "severity": 3,
        "immediate_actions": [
            "Consult a medical doctor for detailed evaluation.",
            "Monitor symptoms for worsening.",
            "Keep the affected area clean and rest.",
        ],
        "refer_to_doctor": True,
        "refer_reason": "Triage report generation formatting error. Recommended safe clinical default.",
        "followup_days": 1,
    }

    if not raw_text:
        return fallback_dict

    # Attempt standard JSON load
    try:
        data = json.loads(raw_text.strip())
        if _is_valid_report(data):
            return data
    except json.JSONDecodeError:
        pass

    # Regex search for JSON object if model outputs wrapper text
    try:
        match = re.search(r"\{.*\}", raw_text, re.DOTALL)
        if match:
            data = json.loads(match.group(0))
            if _is_valid_report(data):
                return data
    except Exception:
        pass

    # Manual field recovery using regular expressions
    try:
        recovered = {}

        cond_match = re.search(r'"likely_condition"\s*:\s*"([^"]+)"', raw_text)
        recovered["likely_condition"] = (
            cond_match.group(1) if cond_match else fallback_dict["likely_condition"]
        )

        sev_match = re.search(r'"severity"\s*:\s*(\d)', raw_text)
        recovered["severity"] = (
            int(sev_match.group(1)) if sev_match else fallback_dict["severity"]
        )

        # Immediate actions list recovery
        actions_match = re.search(
            r'"immediate_actions"\s*:\s*\[(.*?)\]', raw_text, re.DOTALL
        )
        if actions_match:
            actions_text = actions_match.group(1)
            actions = [
                item.strip().strip('"')
                for item in actions_text.split(",")
                if item.strip()
            ]
            recovered["immediate_actions"] = actions
        else:
            recovered["immediate_actions"] = fallback_dict["immediate_actions"]

        refer_match = re.search(
            r'"refer_to_doctor"\s*:\s*(true|false)', raw_text, re.IGNORECASE
        )
        recovered["refer_to_doctor"] = (
            refer_match.group(1).lower() == "true"
            if refer_match
            else fallback_dict["refer_to_doctor"]
        )

        reason_match = re.search(r'"refer_reason"\s*:\s*"([^"]*)"', raw_text)
        recovered["refer_reason"] = reason_match.group(1) if reason_match else ""

        followup_match = re.search(r'"followup_days"\s*:\s*(\d+)', raw_text)
        recovered["followup_days"] = (
            int(followup_match.group(1))
            if followup_match
            else fallback_dict["followup_days"]
        )

        if _is_valid_report(recovered):
            return recovered
    except Exception:
        pass

    return fallback_dict


def _is_valid_report(data: dict) -> bool:
    """Verifies that the required keys are present and have the correct type."""
    required_keys = [
        "likely_condition",
        "severity",
        "immediate_actions",
        "refer_to_doctor",
        "refer_reason",
        "followup_days",
    ]
    return all(k in data for k in required_keys)


_hindi_tts_model = None
_hindi_tts_tokenizer = None


def synthesize_hindi(text: str) -> bytes:
    """Synthesizes Hindi text to speech. Uses local transformers MMS-TTS or gTTS fallback."""
    global _hindi_tts_model, _hindi_tts_tokenizer
    print(f"[Client] Synthesizing Hindi speech for: '{text[:30]}...'")

    # Try loading the local Facebook MMS-TTS-HIN model locally (offline-compatible)
    try:
        if _hindi_tts_model is None or _hindi_tts_tokenizer is None:
            model_id = "facebook/mms-tts-hin"
            _hindi_tts_tokenizer = AutoTokenizer.from_pretrained(model_id)
            _hindi_tts_model = VitsModel.from_pretrained(model_id)

        inputs = _hindi_tts_tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            output = _hindi_tts_model(**inputs).waveform[0].cpu().numpy()

        # Convert waveform to WAV bytes
        out_buf = io.BytesIO()
        wavfile.write(out_buf, rate=16000, data=output)
        return out_buf.getvalue()

    except Exception as e:
        print(
            f"[Client WARNING] Local Hindi TTS Model failed: {e}. Falling back to gTTS API."
        )
        try:
            tts = gTTS(text=text, lang="hi")
            out_buf = io.BytesIO()
            tts.write_to_fp(out_buf)
            return out_buf.getvalue()
        except Exception as api_err:
            print(f"[Client ERROR] gTTS synthesis failed: {api_err}")
            # Return empty audio bytes on complete failure
            return b""


_kokoro_pipeline = None


def synthesize_english(text: str) -> bytes:
    """Synthesizes English text to speech using local Kokoro model or gTTS fallback."""
    global _kokoro_pipeline
    print(f"[Client] Synthesizing English speech for: '{text[:30]}...'")

    try:
        if _kokoro_pipeline is None:
            _kokoro_pipeline = KPipeline(lang_code="a")

        generator = _kokoro_pipeline(text, voice="en-us-1", speed=1.0)
        audio_pieces = []
        sr = 24000
        for _, _, audio in generator:
            if audio is not None:
                audio_pieces.append(audio)

        if not audio_pieces:
            raise ValueError("Kokoro produced no audio arrays")

        combined_audio = np.concatenate(audio_pieces)

        out_buf = io.BytesIO()
        sf.write(out_buf, combined_audio, sr, format="WAV")
        return out_buf.getvalue()

    except Exception as e:
        print(
            f"[Client WARNING] Local Kokoro TTS model failed: {e}. Falling back to gTTS API."
        )
        try:
            tts = gTTS(text=text, lang="en")
            out_buf = io.BytesIO()
            tts.write_to_fp(out_buf)
            return out_buf.getvalue()
        except Exception as api_err:
            print(f"[Client ERROR] English gTTS synthesis failed: {api_err}")
            return b""