main.py

"""
Speech Analysis API — FastAPI backend for Hugging Face Spaces deployment.

Exposes:
  - POST /evaluate      — score a single pronunciation recording
  - POST /diagnose_profile — generate a clinical diagnostic report via GPT-4o
"""

import json
import os
import io
import sys
import uuid
import time
import tempfile
import threading
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, Form, UploadFile, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from pydub import AudioSegment
from openai import OpenAI, APITimeoutError, APIConnectionError

from score_engine import score_pronunciation
from phoneme_extractor import warmup
from data_store import (
    save_evaluation,
    extract_user_id,
    get_user_dashboard,
    get_user_progress,
    save_user_progress,
    _get_client as _init_supabase_client,
)


# =============================================================================
# GPT-4o Structured Output — Clinical Report Schema
# =============================================================================

class ClinicalReport(BaseModel):
    """Structured clinical diagnostic report for a child's speech profile."""
    clinical_hypothesis: str = Field(
        description="A deeply empathetic, professional explanation in Hebrew "
                    "of the root issue behind the child's pronunciation pattern."
    )
    requires_clinic_intervention: bool = Field(
        description="Whether the child needs in-person therapy with a "
                    "speech-language pathologist, or can improve with home practice."
    )
    intervention_reason: str = Field(
        description="Clear explanation in Hebrew of why the child does or "
                    "does not need external clinical therapy."
    )
    estimated_duration_weeks: int = Field(
        description="Estimated number of weeks to reach age-appropriate "
                    "pronunciation based on standard SLP progress curves."
    )
    offline_exercises: list[str] = Field(
        description="Exactly 3 specific, highly actionable offline games or "
                    "exercises the parents can do at home with the child."
    )
    app_roadmap_summary: str = Field(
        description="How the app will adapt its difficulty levels and word "
                    "selection to guide this specific child's improvement."
    )

VALID_WORDS = {
    # ש (Shin) words
    "shalom", "shemesh", "shir", "shuk", "geshem", "shshshsh", "shaon", "shulchan",
    "shin", "shofar", "mishkafayim", "shatiach", "sharsheret", "galshan", "deshe", "yanshuf", "hipushit",
    "mashrokit", "karish", "marshmelo", "machshev", "dvash", "shtayim", "nachash", "sheva", "mivreshet_shinayim",
    # ק (Kuf) words — competitive alignment (K vs T substitution)
    "kof", "kir", "kubiya", "kalmar", "kova", "kadur", "kelev", "mishkafayim_kuf",
    "akavish", "sakin", "mashrokit_kuf",
    # CV syllables (sh-X / k-X) — bridge between isolated sounds and full words
    "sh_syllable_sha", "sh_syllable_she", "sh_syllable_shi", "sh_syllable_shu",
    "k_syllable_ka", "k_syllable_ke", "k_syllable_ki", "k_syllable_ku",
    # Isolated sounds — for children who struggle with full words
    "k_sound", "t_sound", "sh_sound",
}


# How often the background thread re-runs a tiny LOCAL inference to keep the
# Wav2Vec2 model hot. The server is always-on, but after a day idle the model
# path goes cold (memory paged out / CPU throttled) and the first real request
# paid ~30s. /warmup never runs the model, so it can't prevent this. 240s keeps
# the model warm with negligible cost. Override via env if needed.
SELF_WARM_INTERVAL_SEC = int(os.environ.get("SELF_WARM_INTERVAL_SEC", "240"))


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Pre-load Wav2Vec2 model AND Supabase client before accepting requests.

    Without the eager Supabase init, the first /evaluate paid 5-15s on the
    service-role client's import + TLS handshake (logged as
    "[DATA] Supabase client initialized" on the first request only).
    """
    warmup()

    # Eagerly initialise the Supabase service-role client so the first
    # /evaluate doesn't pay the supabase-py import + TLS handshake cost.
    try:
        _init_supabase_client()
    except Exception as e:
        # Don't block server startup — if Supabase is misconfigured the
        # endpoints will still surface the error per-request as before.
        print(f"[WARMUP] Supabase eager init skipped: {e}")

    # Force a dummy /evaluate-style scoring pass so any lazy librosa /
    # numpy / DTW code paths inside score_pronunciation are materialised
    # before the first real child recording arrives.
    try:
        import numpy as np
        import soundfile as sf
        dummy_wav = os.path.join(tempfile.gettempdir(), "warmup_dummy.wav")
        # 0.5s of low-amplitude pink-ish noise — silent enough to be cheap,
        # loud enough that the engine doesn't abort on an empty-signal path.
        rng = np.random.default_rng(0)
        signal = (rng.standard_normal(8000) * 0.01).astype("float32")
        sf.write(dummy_wav, signal, 16000)
        try:
            score_pronunciation(dummy_wav, "k_sound")
            print("[WARMUP] score_pronunciation dummy pass OK")
        except Exception as e:
            print(f"[WARMUP] dummy scoring pass raised (non-fatal): {e}")
        finally:
            if os.path.exists(dummy_wav):
                os.remove(dummy_wav)
    except Exception as e:
        print(f"[WARMUP] dummy scoring setup failed: {e}")

    # Background self-warm: keep the local Wav2Vec2 model HOT 24/7. The server is
    # always-on (paid, never-sleep), but after a day idle the model path goes cold
    # and the first real request paid ~30s. This daemon thread runs a tiny LOCAL
    # inference every few minutes — Wav2Vec2 forward pass only, NO STT API calls
    # (those cost money) — so the model never goes cold, independent of the app.
    def _self_warm_loop():
        import time as _time
        import numpy as np
        import soundfile as sf
        warm_wav = os.path.join(tempfile.gettempdir(), "selfwarm.wav")
        try:
            rng = np.random.default_rng(1)
            sf.write(warm_wav, (rng.standard_normal(8000) * 0.01).astype("float32"), 16000)
        except Exception as e:
            print(f"[SELF-WARM] setup failed, warmer not started: {e}")
            return
        from phoneme_extractor import shin_vs_samekh
        while True:
            _time.sleep(SELF_WARM_INTERVAL_SEC)
            try:
                shin_vs_samekh(warm_wav)  # local Wav2Vec2 forward pass, no network
                print("[SELF-WARM] model kept hot")
            except Exception as e:
                print(f"[SELF-WARM] tick failed: {e}")

    threading.Thread(target=_self_warm_loop, daemon=True).start()
    print(f"[SELF-WARM] background warmer started (every {SELF_WARM_INTERVAL_SEC}s)")

    yield


app = FastAPI(title="Speech Analysis API", version="1.0.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# =============================================================================
# Warm-up endpoints — used by the Unity client to keep the TLS connection
# alive and prove the Space is ready, without paying the cost of /evaluate.
# =============================================================================

@app.get("/")
async def root():
    """Root endpoint — replaces the 404 the warming pings used to hit."""
    return {"ok": True, "service": "speechkid-api"}


@app.post("/warmup")
async def warmup_endpoint():
    """
    Cheap POST that establishes / refreshes the TLS connection on the same
    code path as /evaluate (POST + JSON body), without running scoring.
    The Unity client pings this at boot and then every ~30s so the first
    real /evaluate doesn't pay the TLS handshake cost.
    """
    return {"ok": True, "warm": True}


def convert_to_wav(input_path: str, output_path: str) -> None:
    """Convert any audio format (m4a, aac, caf, webm, etc.) to 16kHz mono WAV."""
    audio = AudioSegment.from_file(input_path)
    audio = audio.set_frame_rate(16000).set_channels(1).set_sample_width(2)
    audio.export(output_path, format="wav")


@app.post("/evaluate")
async def evaluate(
    file: UploadFile = File(...),
    word: str = Form(...),
    x_supabase_token: str = Header(default=None),
):
    """
    Evaluate a child's pronunciation of a Hebrew word.

    Supports two phoneme types:
    - **ש (Shin) words**: shalom, shemesh, shir, shuk, geshem, shshshsh, shaon, shulchan
    - **ק (Kuf) words**: kof, kir, kubiya, kalmar (detects T-for-K substitution)

    - **file**: Audio recording (WAV, M4A, AAC, CAF, WebM, etc.).
    - **word**: Target word key.

    Returns the full scoring result including diagnosis and feedback.
    """
    if word not in VALID_WORDS:
        raise HTTPException(
            status_code=422,
            detail=f"Unknown word '{word}'. Must be one of: {', '.join(sorted(VALID_WORDS))}",
        )

    # Verify Supabase auth token and extract user_id
    user_id = extract_user_id(x_supabase_token)

    evaluation_id = str(uuid.uuid4())
    uid = evaluation_id.replace("-", "")
    suffix = os.path.splitext(file.filename or "recording.bin")[1] or ".bin"
    raw_path = os.path.join(tempfile.gettempdir(), f"{uid}_raw{suffix}")
    wav_path = os.path.join(tempfile.gettempdir(), f"{uid}.wav")

    try:
        # Save the uploaded file with its original extension
        contents = await file.read()
        with open(raw_path, "wb") as f:
            f.write(contents)

        # Convert to 16kHz mono WAV (handles m4a, aac, caf, webm, ogg, etc.)
        convert_to_wav(raw_path, wav_path)

        # Capture server logs during scoring
        log_buffer = io.StringIO()
        old_stdout = sys.stdout
        sys.stdout = log_buffer

        start_time = time.time()
        result = score_pronunciation(wav_path, word)
        processing_ms = int((time.time() - start_time) * 1000)

        sys.stdout = old_stdout
        server_logs = log_buffer.getvalue()

        # Print captured logs to actual stdout so they appear in HF Spaces logs
        print(server_logs, end="")

        # Read WAV into memory before files get deleted in finally block
        with open(wav_path, "rb") as f:
            wav_data = f.read()

        # Save recording + result to Supabase (background thread, non-blocking)
        save_evaluation(evaluation_id, word, result, wav_data, server_logs, processing_ms, user_id)

        return result

    except Exception:
        # Restore stdout if scoring crashed mid-capture
        sys.stdout = sys.__stdout__
        raise

    finally:
        for p in (raw_path, wav_path):
            if os.path.exists(p):
                os.remove(p)


# =============================================================================
# Dashboard — User Stats & Recent Activity
# =============================================================================

@app.get("/api/v1/dashboard")
async def dashboard(x_supabase_token: str = Header(default=None)):
    """
    Return the logged-in user's practice stats and recent evaluations.

    Authenticates via X-Supabase-Token header. Queries are scoped to the
    user via RLS (uses anon key + user token, not service_role).
    """
    if not x_supabase_token:
        raise HTTPException(status_code=401, detail="Missing X-Supabase-Token header")

    # Verify the token is valid
    user_id = extract_user_id(x_supabase_token)
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid or expired token")

    # Query user's data (RLS-scoped via user token + anon key)
    result = get_user_dashboard(x_supabase_token)
    if result is None:
        raise HTTPException(status_code=500, detail="Dashboard query failed — SUPABASE_ANON_KEY may not be configured")

    return result


# =============================================================================
# Progress — Level Completion & Unlocks
# =============================================================================

class ProgressPayload(BaseModel):
    level: int = Field(ge=1)
    stars: int = Field(ge=0)
    passed: bool
    word_count: int = Field(ge=0)


@app.get("/api/v1/progress")
async def get_progress(x_supabase_token: str = Header(default=None)):
    """Return the user's unlocked level and per-level star breakdown."""
    if not x_supabase_token:
        raise HTTPException(status_code=401, detail="Missing X-Supabase-Token header")

    user_id = extract_user_id(x_supabase_token)
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid or expired token")

    result = get_user_progress(x_supabase_token)
    if result is None:
        raise HTTPException(
            status_code=500,
            detail="Progress query failed — SUPABASE_ANON_KEY may not be configured",
        )
    return result


@app.post("/api/v1/progress")
async def post_progress(
    payload: ProgressPayload,
    x_supabase_token: str = Header(default=None),
):
    """Upsert level completion with merge semantics (max stars, sticky passed)."""
    if not x_supabase_token:
        raise HTTPException(status_code=401, detail="Missing X-Supabase-Token header")

    user_id = extract_user_id(x_supabase_token)
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid or expired token")

    ok = save_user_progress(
        x_supabase_token,
        level=payload.level,
        stars=payload.stars,
        passed=payload.passed,
        word_count=payload.word_count,
    )
    if not ok:
        raise HTTPException(
            status_code=500,
            detail="Progress save failed — SUPABASE_ANON_KEY may not be configured",
        )
    return {"ok": True}


# =============================================================================
# Diagnostic Profile — GPT-4o Clinical Reasoning
# =============================================================================

SYSTEM_PROMPT = """\
You are a Top-Tier Israeli Pediatric Speech-Language Pathologist (קלינאית תקשורת מומחית). \
Your task is to analyze an intake form for a child and provide a highly accurate, empathetic, \
and professional diagnostic report in fluent Hebrew.

CRITICAL CLINICAL REASONING INSTRUCTIONS:
Before generating the structured JSON output, you MUST internally follow these logical steps (Chain of Thought):
1. Analyze Age vs. Norms: Compare the child's age to Hebrew acquisition norms. \
   - Basic sounds (B, P, M, N, T, D, K, G, Ch, R, L) should be mastered by age 4. \
   - Hissing sounds (S, Z, Ts, Sh) are typically mastered by age 6. \
   - Note: If a 4-year-old struggles with 'SH', it is developmental. If a 7-year-old does, it requires intervention.
2. Differentiate Articulation vs. Phonology: \
   - Articulation (Phonetic): Distorting a specific sound (e.g., lateral lisp on 'S' or 'SH', guttural 'R'). \
   - Phonological: Changing word structures. E.g., Fronting (K->T), Cluster Reduction ("Psanter" -> "Santer"), \
Dropping ending consonants. Phonological errors past age 4-5 are highly significant.
3. Check Red Flags: If any medical/organic red flags (cleft palate, hearing loss, regression, drooling) \
are marked TRUE, you MUST recommend immediate medical consultation.
4. Assess Intelligibility & Frustration: If intelligibility is low or frustration is high, the tone must be \
highly empathetic, and the app plan must start with high "forgiveness" (lenient grading) and auditory \
discrimination before verbal production.

Generate the exact JSON structure required by the Pydantic schema based on this reasoning. \
The tone must be professional, warm, reassuring, and clearly state that this is an AI-assisted \
roadmap and not a replacement for a formal clinical evaluation."""


class IntakeForm(BaseModel):
    """Intake questionnaire submitted by the parent."""
    child_age_years: int = Field(description="Child's age in years")
    child_age_months: int = Field(default=0, description="Additional months")
    target_sounds: list[str] = Field(
        description="Sounds the child struggles with, e.g. ['ש', 'ר']"
    )
    questionnaire: dict = Field(
        description="Free-form questionnaire answers from the parent"
    )


@app.post("/diagnose_profile", response_model=ClinicalReport)
async def diagnose_profile(form: IntakeForm):
    """
    Generate a clinical diagnostic report based on a parent's intake form.

    Uses GPT-4o with Structured Outputs to guarantee valid JSON matching
    the ClinicalReport schema. All output is in Hebrew.

    - **child_age_years**: Child's age in years.
    - **child_age_months**: Additional months (optional, default 0).
    - **target_sounds**: List of sounds the child struggles with.
    - **questionnaire**: Dict of parent-provided intake answers.

    Returns a ClinicalReport with hypothesis, exercises, and recommendations.
    """
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise HTTPException(
            status_code=500,
            detail="OPENAI_API_KEY environment variable is not set.",
        )

    age_display = f"{form.child_age_years} שנים"
    if form.child_age_months:
        age_display += f" ו-{form.child_age_months} חודשים"

    user_prompt = (
        f"נתוני הילד/ה:\n"
        f"גיל: {age_display}\n"
        f"צלילים בעייתיים: {', '.join(form.target_sounds)}\n\n"
        f"תשובות שאלון ההורים:\n"
        f"{json.dumps(form.questionnaire, ensure_ascii=False, indent=2)}"
    )

    client = OpenAI(api_key=api_key, timeout=60.0)

    try:
        completion = client.beta.chat.completions.parse(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            response_format=ClinicalReport,
        )
    except (APITimeoutError, APIConnectionError) as e:
        raise HTTPException(
            status_code=504,
            detail=f"OpenAI API timeout or connection error: {e}",
        )
    except Exception as e:
        raise HTTPException(
            status_code=502,
            detail=f"OpenAI API error: {e}",
        )

    report = completion.choices[0].message.parsed
    if report is None:
        raise HTTPException(
            status_code=502,
            detail="OpenAI returned a response but structured parsing failed.",
        )

    return report