"""Vaccination repository for database operations on vaccinations."""
from typing import List, Optional

from bson import ObjectId
from pymongo import ReturnDocument

from app.db import Database
from app.models.vaccination import Vaccination

COL = "vaccinations"

def create(vaccination: Vaccination) -> Vaccination:
    """Create a new vaccination record in the database."""
    vaccination_to_insert = vaccination.to_bson()
    res = Database.get_db().get_collection(COL).insert_one(vaccination_to_insert)
    return Vaccination.from_bson({**vaccination_to_insert, "_id": res.inserted_id})

def find_by_id(vaccination_id: str) -> Optional[Vaccination]:
    """Find a vaccination record by its ID."""
    if not ObjectId.is_valid(vaccination_id):
        return None
    doc = Database.get_db().get_collection(COL).find_one({"_id": ObjectId(vaccination_id)})
    return Vaccination.from_bson(doc) if doc else None

def find_by_patient_id(patient_id: str) -> List[Vaccination]:
    """Find all vaccination records for a patient."""
    if not ObjectId.is_valid(patient_id):
        return []
    docs = Database.get_db().get_collection(COL).find({"patientId": ObjectId(patient_id)})
    return [Vaccination.from_bson(doc) for doc in docs]

def update(vaccination_id: str, updates: dict) -> Optional[Vaccination]:
    """Update a vaccination record."""
    if not ObjectId.is_valid(vaccination_id):
        return None
    db = Database.get_db()
    existing = db.get_collection(COL).find_one({"_id": ObjectId(vaccination_id)})
    if not existing:
        return None

    # Merge existing document with updates, validate via model
    merged = {**existing, **updates}
    candidate = Vaccination.from_bson(merged)
    to_set = candidate.to_bson()
    to_set.pop("_id", None)

    doc = db.get_collection(COL).find_one_and_update(
        {"_id": ObjectId(vaccination_id)},
        {"$set": to_set},
        return_document=ReturnDocument.AFTER,
    )
    return Vaccination.from_bson(doc) if doc else None
