backend/models/jarvis.py

"""Persistence layer for the Jarvis explainability + feedback + chat loops.

Four tables:

* ``jarvis_explanations`` — one row per analysis the reasoner produces
  for an alert-level reading (warning / critical). The full structured
  payload (top features, recommendations, RAG passage ids, generated
  clinical note) is stored as JSON so we can replay an explanation
  weeks later without rerunning the model. This is the audit trail
  the ANVISA-compliant deployment depends on.

* ``jarvis_feedback`` — clinician-supplied corrections that close the
  HIL loop. A correction always points back to an explanation row so
  we can retrain on the model's actual outputs rather than on raw
  sensor data.

* ``jarvis_conversations`` — one row per chat thread between a
  clinician and the local Llama LLM. Optionally linked to a patient
  so the chat carries live BioPatch context into every turn.

* ``jarvis_messages`` — ordered user/assistant turns inside a
  conversation, plus the RAG passage ids that grounded each
  assistant response.
"""

from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING

from sqlalchemy import JSON, DateTime, Float, ForeignKey, Integer, String, Text, func
from sqlalchemy.orm import Mapped, mapped_column, relationship

from core.database import Base

if TYPE_CHECKING:
    from .patient import Patient
    from .user import User
    from .vitals import VitalsHistory


class JarvisExplanation(Base):
    __tablename__ = "jarvis_explanations"

    id: Mapped[int] = mapped_column(primary_key=True)
    patient_id: Mapped[int] = mapped_column(
        ForeignKey("patients.id", ondelete="CASCADE"), index=True, nullable=False
    )
    vitals_id: Mapped[int | None] = mapped_column(
        ForeignKey("vitals_history.id", ondelete="SET NULL"), index=True, nullable=True
    )
    severity: Mapped[str] = mapped_column(String(16), index=True, nullable=False)
    score: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)
    confidence: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)
    lang: Mapped[str] = mapped_column(String(4), default="en", nullable=False)
    bacterial_signature: Mapped[str | None] = mapped_column(String(128), nullable=True)
    tissue_state: Mapped[str | None] = mapped_column(String(32), nullable=True)
    note_headline: Mapped[str | None] = mapped_column(String(255), nullable=True)
    note_body: Mapped[str | None] = mapped_column(Text, nullable=True)
    reasoning: Mapped[dict] = mapped_column(JSON, nullable=False)
    citations: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
    generator: Mapped[str] = mapped_column(String(32), default="template-v1", nullable=False)
    model_version: Mapped[str | None] = mapped_column(String(64), nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), index=True, nullable=False
    )

    patient: Mapped["Patient"] = relationship("Patient")
    vitals: Mapped["VitalsHistory | None"] = relationship("VitalsHistory")
    feedback: Mapped[list["JarvisFeedback"]] = relationship(
        "JarvisFeedback", back_populates="explanation", cascade="all, delete-orphan"
    )

    def __repr__(self) -> str:
        return (
            f"<JarvisExplanation id={self.id} patient_id={self.patient_id} "
            f"severity={self.severity!r} score={self.score}>"
        )


class JarvisFeedback(Base):
    """Clinician correction of a Jarvis explanation (the HIL loop)."""

    __tablename__ = "jarvis_feedback"

    id: Mapped[int] = mapped_column(primary_key=True)
    explanation_id: Mapped[int] = mapped_column(
        ForeignKey("jarvis_explanations.id", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    author_id: Mapped[int] = mapped_column(ForeignKey("users.id"), nullable=False)
    correct_severity: Mapped[str] = mapped_column(String(16), nullable=False)
    correction_kind: Mapped[str] = mapped_column(
        String(32), default="confirmation", nullable=False
    )  # confirmation | false_positive | false_negative | nuance
    clinician_notes: Mapped[str | None] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )

    explanation: Mapped["JarvisExplanation"] = relationship("JarvisExplanation", back_populates="feedback")
    author: Mapped["User"] = relationship("User")

    def __repr__(self) -> str:
        return (
            f"<JarvisFeedback id={self.id} explanation_id={self.explanation_id} "
            f"kind={self.correction_kind!r}>"
        )


class JarvisConversation(Base):
    """A chat thread between a clinician and the local Jarvis LLM."""

    __tablename__ = "jarvis_conversations"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("users.id", ondelete="CASCADE"), index=True, nullable=False
    )
    patient_id: Mapped[int | None] = mapped_column(
        ForeignKey("patients.id", ondelete="SET NULL"), index=True, nullable=True
    )
    title: Mapped[str] = mapped_column(String(255), default="New conversation", nullable=False)
    lang: Mapped[str] = mapped_column(String(4), default="en", nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), index=True, nullable=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        onupdate=func.now(),
        nullable=False,
    )

    user: Mapped["User"] = relationship("User")
    patient: Mapped["Patient | None"] = relationship("Patient")
    messages: Mapped[list["JarvisMessage"]] = relationship(
        "JarvisMessage",
        back_populates="conversation",
        cascade="all, delete-orphan",
        order_by="JarvisMessage.created_at",
    )

    def __repr__(self) -> str:
        return (
            f"<JarvisConversation id={self.id} user_id={self.user_id} "
            f"patient_id={self.patient_id} title={self.title!r}>"
        )


class JarvisMessage(Base):
    """A single user or assistant turn inside a Jarvis chat."""

    __tablename__ = "jarvis_messages"

    id: Mapped[int] = mapped_column(primary_key=True)
    conversation_id: Mapped[int] = mapped_column(
        ForeignKey("jarvis_conversations.id", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    role: Mapped[str] = mapped_column(String(16), nullable=False)  # user | assistant | system
    content: Mapped[str] = mapped_column(Text, nullable=False)
    citations: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
    rag_passage_ids: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
    actions: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
    token_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
    finished: Mapped[bool] = mapped_column(default=True, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), index=True, nullable=False
    )

    conversation: Mapped["JarvisConversation"] = relationship(
        "JarvisConversation", back_populates="messages"
    )

    def __repr__(self) -> str:
        return (
            f"<JarvisMessage id={self.id} conv={self.conversation_id} "
            f"role={self.role!r}>"
        )


class JarvisPreference(Base):
    """Per-user key/value preferences that Jarvis adapts to over time.

    Keys are free-form but the standard set is:

    * ``lang``                — clinician-preferred response language.
    * ``last_patient_id``     — the patient most recently focused.
    * ``voice_enabled``       — whether to auto-speak Jarvis replies.
    * ``wake_greeting``       — whether to play the welcome on session.
    * ``preferred_body_view`` — front / back / xray / thermal.
    """

    __tablename__ = "jarvis_preferences"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("users.id", ondelete="CASCADE"), index=True, nullable=False
    )
    key: Mapped[str] = mapped_column(String(64), nullable=False)
    value: Mapped[str] = mapped_column(Text, nullable=False)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        onupdate=func.now(),
        nullable=False,
    )

    user: Mapped["User"] = relationship("User")

    def __repr__(self) -> str:
        return f"<JarvisPreference user_id={self.user_id} key={self.key!r}>"


class JarvisCommandLog(Base):
    """Audit trail for voice / text commands and their outcomes.

    The reasoner reads this table to bias future routing — e.g. if
    the same clinician repeatedly says "open les patients" and the
    rule router misses it 3 times in a row, an admin can promote
    that phrasing to the multilingual alias table.
    """

    __tablename__ = "jarvis_command_log"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("users.id", ondelete="CASCADE"), index=True, nullable=False
    )
    transcript: Mapped[str] = mapped_column(Text, nullable=False)
    transport: Mapped[str] = mapped_column(String(16), default="text", nullable=False)  # text | voice
    lang: Mapped[str] = mapped_column(String(4), default="en", nullable=False)
    action_name: Mapped[str | None] = mapped_column(String(64), nullable=True)
    action_params: Mapped[dict] = mapped_column(JSON, nullable=False, default=dict)
    confidence: Mapped[float | None] = mapped_column(Float, nullable=True)
    outcome: Mapped[str] = mapped_column(
        String(16), default="pending", nullable=False
    )  # pending | ok | failed | corrected
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), index=True, nullable=False
    )

    user: Mapped["User"] = relationship("User")

    def __repr__(self) -> str:
        return (
            f"<JarvisCommandLog id={self.id} user={self.user_id} "
            f"action={self.action_name!r} outcome={self.outcome!r}>"
        )


class ConfirmedInfection(Base):
    """Doctor-validated infection record — the institutional memory.

    Created by a wound specialist (or admin) when they confirm that a
    Jarvis alert corresponded to a real surgical-site infection. The
    record snapshots the wound vitals at confirmation, links back to
    the originating Jarvis explanation, and captures the clinical
    outcome (organism, treatment, time to resolution).

    The case-retrieval engine queries this table at every new alert,
    surfaces the K most similar historical cases as empirical
    evidence, and updates the BioPatch institutional memory in real
    time without ever auto-tuning thresholds.
    """

    __tablename__ = "jarvis_confirmed_infections"

    id: Mapped[int] = mapped_column(primary_key=True)
    patient_id: Mapped[int] = mapped_column(
        ForeignKey("patients.id", ondelete="CASCADE"), index=True, nullable=False
    )
    explanation_id: Mapped[int | None] = mapped_column(
        ForeignKey("jarvis_explanations.id", ondelete="SET NULL"),
        index=True,
        nullable=True,
    )
    vitals_id: Mapped[int | None] = mapped_column(
        ForeignKey("vitals_history.id", ondelete="SET NULL"), nullable=True
    )
    confirmed_by_id: Mapped[int] = mapped_column(
        ForeignKey("users.id"), nullable=False
    )

    # Snapshot of the wound + patient at confirmation
    temperature_c: Mapped[float] = mapped_column(Float, nullable=False)
    ph: Mapped[float] = mapped_column(Float, nullable=False)
    humidity: Mapped[float] = mapped_column(Float, nullable=False)
    severity_at_confirmation: Mapped[str] = mapped_column(String(16), nullable=False)
    score_at_confirmation: Mapped[float] = mapped_column(Float, default=0.0, nullable=False)
    age_at_confirmation: Mapped[int] = mapped_column(Integer, nullable=False)
    diabetes: Mapped[bool] = mapped_column(default=False, nullable=False)
    smoker: Mapped[bool] = mapped_column(default=False, nullable=False)
    sex: Mapped[str] = mapped_column(String(1), default="X", nullable=False)
    surgery_type: Mapped[str] = mapped_column(String(128), nullable=False)
    days_post_op: Mapped[int | None] = mapped_column(Integer, nullable=True)

    # Trajectory (free-form snapshot of recent vitals at confirmation time)
    vitals_trajectory: Mapped[list] = mapped_column(JSON, nullable=False, default=list)

    # Clinical outcome
    organism: Mapped[str | None] = mapped_column(String(128), nullable=True)
    organism_confidence: Mapped[float | None] = mapped_column(Float, nullable=True)
    treatment: Mapped[str | None] = mapped_column(Text, nullable=True)
    resolution_days: Mapped[int | None] = mapped_column(Integer, nullable=True)
    outcome: Mapped[str] = mapped_column(
        String(32), default="ongoing", nullable=False
    )  # ongoing | resolved | escalated | deceased | unresolved

    # Doctor's free-form notes
    clinical_notes: Mapped[str | None] = mapped_column(Text, nullable=True)
    tags: Mapped[list] = mapped_column(JSON, nullable=False, default=list)

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), index=True, nullable=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        onupdate=func.now(),
        nullable=False,
    )

    patient: Mapped["Patient"] = relationship("Patient")
    confirmed_by: Mapped["User"] = relationship("User")

    def __repr__(self) -> str:
        return (
            f"<ConfirmedInfection id={self.id} patient_id={self.patient_id} "
            f"organism={self.organism!r} outcome={self.outcome!r}>"
        )


class DoctorCorrection(Base):
    """Rich correction record from the wound specialist.

    Goes beyond the binary :class:`JarvisFeedback` (false positive /
    confirmation) to capture *what* Jarvis got wrong on multiple axes:
    severity, bacterial signature, tissue state, recommended actions.
    Each axis is stored separately so the learning analytics can show
    Jarvis's per-dimension agreement rate with the specialist.
    """

    __tablename__ = "jarvis_doctor_corrections"

    id: Mapped[int] = mapped_column(primary_key=True)
    explanation_id: Mapped[int] = mapped_column(
        ForeignKey("jarvis_explanations.id", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    doctor_id: Mapped[int] = mapped_column(ForeignKey("users.id"), nullable=False)

    # Per-axis corrections — None = doctor agreed with Jarvis on that axis
    corrected_severity: Mapped[str | None] = mapped_column(String(16), nullable=True)
    corrected_signature: Mapped[str | None] = mapped_column(String(128), nullable=True)
    corrected_tissue_state: Mapped[str | None] = mapped_column(String(32), nullable=True)

    rationale: Mapped[str | None] = mapped_column(Text, nullable=True)
    teaching_point: Mapped[str | None] = mapped_column(Text, nullable=True)

    # Whether this correction was promoted into a ConfirmedInfection
    promoted_to_case_id: Mapped[int | None] = mapped_column(
        ForeignKey("jarvis_confirmed_infections.id", ondelete="SET NULL"),
        nullable=True,
    )

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), nullable=False
    )

    doctor: Mapped["User"] = relationship("User")

    def __repr__(self) -> str:
        return (
            f"<DoctorCorrection id={self.id} explanation_id={self.explanation_id} "
            f"by={self.doctor_id}>"
        )