normogen/backend/src/models/medication.rs

use super::health_data::EncryptedField;
use futures::stream::StreamExt;
use mongodb::bson::{doc, oid::ObjectId, DateTime};
use mongodb::Collection;
use serde::{Deserialize, Serialize};

// ============================================================================
// PILL IDENTIFICATION (Phase 2.8)
// ============================================================================

/// Physical pill identification (optional)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PillIdentification {
    /// Size of the pill (optional)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub size: Option<PillSize>,

    /// Shape of the pill (optional)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub shape: Option<PillShape>,

    /// Color of the pill (optional)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub color: Option<PillColor>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PillSize {
    Tiny,       // < 5mm
    Small,      // 5-10mm
    Medium,     // 10-15mm
    Large,      // 15-20mm
    ExtraLarge, // > 20mm
    #[serde(rename = "custom")]
    Custom(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PillShape {
    Round,
    Oval,
    Oblong,
    Capsule,
    Tablet,
    Square,
    Rectangular,
    Triangular,
    Diamond,
    Hexagonal,
    Octagonal,
    #[serde(rename = "custom")]
    Custom(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PillColor {
    White,
    OffWhite,
    Yellow,
    Orange,
    Red,
    Pink,
    Purple,
    Blue,
    Green,
    Brown,
    Black,
    Gray,
    Clear,
    #[serde(rename = "multi-colored")]
    MultiColored,
    #[serde(rename = "custom")]
    Custom(String),
}

// ============================================================================
// ADHERENCE STATISTICS (Phase 2.7)
// ============================================================================

/// Adherence statistics calculated for a medication
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdherenceStats {
    pub medication_id: String,
    pub total_doses: i64,
    pub scheduled_doses: i64,
    pub taken_doses: i64,
    pub missed_doses: i64,
    pub adherence_rate: f64,
    pub period_days: i64,
}

// ============================================================================
// MEDICATION MODEL (Existing + Phase 2.8 updates)
// ============================================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Medication {
    #[serde(rename = "_id", skip_serializing_if = "Option::is_none")]
    pub id: Option<ObjectId>,
    #[serde(rename = "medicationId")]
    pub medication_id: String,
    #[serde(rename = "userId")]
    pub user_id: String,
    #[serde(rename = "profileId")]
    pub profile_id: String,
    #[serde(rename = "medicationData")]
    pub medication_data: EncryptedField,
    #[serde(rename = "reminders")]
    pub reminders: Vec<MedicationReminder>,
    #[serde(rename = "createdAt")]
    pub created_at: DateTime,
    #[serde(rename = "updatedAt")]
    pub updated_at: DateTime,

    /// Physical pill identification (Phase 2.8 - optional)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pill_identification: Option<PillIdentification>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MedicationReminder {
    #[serde(rename = "reminderId")]
    pub reminder_id: String,
    #[serde(rename = "scheduledTime")]
    pub scheduled_time: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MedicationDose {
    #[serde(rename = "_id", skip_serializing_if = "Option::is_none")]
    pub id: Option<ObjectId>,
    #[serde(rename = "medicationId")]
    pub medication_id: String,
    #[serde(rename = "userId")]
    pub user_id: String,
    #[serde(rename = "loggedAt")]
    pub logged_at: DateTime,
    #[serde(rename = "scheduledTime")]
    pub scheduled_time: Option<String>,
    #[serde(rename = "taken")]
    pub taken: bool,
    #[serde(rename = "notes")]
    pub notes: Option<String>,
}

// ============================================================================
// REQUEST TYPES FOR API (Updated for Phase 2.8)
// ============================================================================

#[derive(Debug, Deserialize)]
pub struct CreateMedicationRequest {
    pub name: String,
    pub dosage: String,
    pub frequency: String,
    pub route: String,
    pub reason: Option<String>,
    pub instructions: Option<String>,
    pub side_effects: Option<Vec<String>>,
    pub prescribed_by: Option<String>,
    pub prescribed_date: Option<String>,
    pub start_date: Option<String>,
    pub end_date: Option<String>,
    pub notes: Option<String>,
    pub tags: Option<Vec<String>>,
    pub reminder_times: Option<Vec<String>>,
    pub profile_id: String,

    /// Pill identification (Phase 2.8 - optional)
    #[serde(rename = "pill_identification")]
    pub pill_identification: Option<PillIdentification>,
}

#[derive(Debug, Deserialize)]
pub struct UpdateMedicationRequest {
    pub name: Option<String>,
    pub dosage: Option<String>,
    pub frequency: Option<String>,
    pub route: Option<String>,
    pub reason: Option<String>,
    pub instructions: Option<String>,
    pub side_effects: Option<Vec<String>>,
    pub prescribed_by: Option<String>,
    pub prescribed_date: Option<String>,
    pub start_date: Option<String>,
    pub end_date: Option<String>,
    pub notes: Option<String>,
    pub tags: Option<Vec<String>>,
    pub reminder_times: Option<Vec<String>>,

    /// Pill identification (Phase 2.8 - optional)
    #[serde(rename = "pill_identification")]
    pub pill_identification: Option<PillIdentification>,
}

#[derive(Debug, Deserialize)]
pub struct LogDoseRequest {
    pub taken: Option<bool>,
    pub scheduled_time: Option<String>,
    pub notes: Option<String>,
}

// ============================================================================
// REPOSITORY
// ============================================================================

/// Repository for Medication operations
#[derive(Clone)]
pub struct MedicationRepository {
    collection: Collection<Medication>,
}

impl MedicationRepository {
    pub fn new(collection: Collection<Medication>) -> Self {
        Self { collection }
    }

    pub async fn create(
        &self,
        medication: Medication,
    ) -> Result<Medication, Box<dyn std::error::Error>> {
        let _result = self.collection.insert_one(medication.clone(), None).await?;
        Ok(medication)
    }

    pub async fn find_by_user(
        &self,
        user_id: &str,
    ) -> Result<Vec<Medication>, Box<dyn std::error::Error>> {
        let filter = doc! { "userId": user_id };
        let mut cursor = self.collection.find(filter, None).await?;
        let mut medications = Vec::new();
        while let Some(medication) = cursor.next().await {
            medications.push(medication?);
        }
        Ok(medications)
    }

    pub async fn find_by_user_and_profile(
        &self,
        user_id: &str,
        profile_id: &str,
    ) -> Result<Vec<Medication>, Box<dyn std::error::Error>> {
        let filter = doc! {
            "userId": user_id,
            "profileId": profile_id
        };
        let mut cursor = self.collection.find(filter, None).await?;
        let mut medications = Vec::new();
        while let Some(medication) = cursor.next().await {
            medications.push(medication?);
        }
        Ok(medications)
    }

    pub async fn find_by_id(
        &self,
        id: &ObjectId,
    ) -> Result<Option<Medication>, Box<dyn std::error::Error>> {
        let filter = doc! { "_id": id };
        let medication = self.collection.find_one(filter, None).await?;
        Ok(medication)
    }

    pub async fn update(
        &self,
        id: &ObjectId,
        updates: UpdateMedicationRequest,
    ) -> Result<Option<Medication>, Box<dyn std::error::Error>> {
        let mut update_doc = doc! {};

        if let Some(name) = updates.name {
            update_doc.insert("medicationData.name", name);
        }
        if let Some(dosage) = updates.dosage {
            update_doc.insert("medicationData.dosage", dosage);
        }
        if let Some(frequency) = updates.frequency {
            update_doc.insert("medicationData.frequency", frequency);
        }
        if let Some(route) = updates.route {
            update_doc.insert("medicationData.route", route);
        }
        if let Some(reason) = updates.reason {
            update_doc.insert("medicationData.reason", reason);
        }
        if let Some(instructions) = updates.instructions {
            update_doc.insert("medicationData.instructions", instructions);
        }
        if let Some(side_effects) = updates.side_effects {
            update_doc.insert("medicationData.sideEffects", side_effects);
        }
        if let Some(prescribed_by) = updates.prescribed_by {
            update_doc.insert("medicationData.prescribedBy", prescribed_by);
        }
        if let Some(prescribed_date) = updates.prescribed_date {
            update_doc.insert("medicationData.prescribedDate", prescribed_date);
        }
        if let Some(start_date) = updates.start_date {
            update_doc.insert("medicationData.startDate", start_date);
        }
        if let Some(end_date) = updates.end_date {
            update_doc.insert("medicationData.endDate", end_date);
        }
        if let Some(notes) = updates.notes {
            update_doc.insert("medicationData.notes", notes);
        }
        if let Some(tags) = updates.tags {
            update_doc.insert("medicationData.tags", tags);
        }
        if let Some(reminder_times) = updates.reminder_times {
            update_doc.insert("reminderTimes", reminder_times);
        }
        if let Some(pill_identification) = updates.pill_identification {
            if let Ok(pill_doc) = mongodb::bson::to_document(&pill_identification) {
                update_doc.insert("pillIdentification", pill_doc);
            }
        }

        update_doc.insert("updatedAt", mongodb::bson::DateTime::now());

        let filter = doc! { "_id": id };
        let medication = self
            .collection
            .find_one_and_update(filter, doc! { "$set": update_doc }, None)
            .await?;
        Ok(medication)
    }

    pub async fn delete(&self, id: &ObjectId) -> Result<bool, Box<dyn std::error::Error>> {
        let filter = doc! { "_id": id };
        let result = self.collection.delete_one(filter, None).await?;
        Ok(result.deleted_count > 0)
    }

    pub async fn calculate_adherence(
        &self,
        medication_id: &str,
        days: i64,
    ) -> Result<AdherenceStats, Box<dyn std::error::Error>> {
        // For now, return a placeholder adherence calculation
        // In a full implementation, this would query the medication_doses collection
        Ok(AdherenceStats {
            medication_id: medication_id.to_string(),
            total_doses: 0,
            scheduled_doses: 0,
            taken_doses: 0,
            missed_doses: 0,
            adherence_rate: 100.0,
            period_days: days,
        })
    }
}