CIA/e-voting-system/backend/scripts/seed_db.py

"""
Script de réinitialisation et de peuplement de la base de données pour tests.

Usage (depuis la racine du projet ou à l'intérieur du conteneur):
    python -m backend.scripts.seed_db

Ce script supprime/crée les tables et insère des électeurs, élections,
candidats et votes d'exemple selon la demande de l'utilisateur.
"""
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
from ..database import engine, SessionLocal
from ..models import Base, Voter, Election, Candidate, Vote
from ..auth import hash_password


def reset_db():
    print("Resetting database: dropping and recreating all tables...")
    Base.metadata.drop_all(bind=engine)
    Base.metadata.create_all(bind=engine)


def seed():
    print("Seeding database with sample voters, elections, candidates and votes...")
    db: Session = SessionLocal()
    try:
        now = datetime.utcnow()

        # Create voters with the requested CNI values
        voters = []
        for cid in ["1234", "12345", "123456"]:
            v = Voter(
                email=f"user{cid}@example.com",
                password_hash=hash_password("Password123"),
                first_name=f"User{cid}",
                last_name="Seed",
                citizen_id=cid
            )
            db.add(v)
            voters.append(v)
        db.commit()
        for v in voters:
            db.refresh(v)

        # Create a few elections: past, present, future
        elections = []

        # Past elections (ended 1 month ago)
        for i in range(1, 4):
            e = Election(
                name=f"Past Election {i}",
                description="Election seed (past)",
                start_date=now - timedelta(days=60 + i),
                end_date=now - timedelta(days=30 + i),
                is_active=False,
                results_published=True
            )
            db.add(e)
            elections.append(e)

        # Current election (ends in ~3 months)
        current = Election(
            name="Ongoing Election",
            description="Election seed (present)",
            start_date=now - timedelta(days=1),
            end_date=now + timedelta(days=90),
            is_active=True,
            results_published=False
        )
        db.add(current)
        elections.append(current)

        # Future elections
        for i in range(1, 4):
            e = Election(
                name=f"Future Election {i}",
                description="Election seed (future)",
                start_date=now + timedelta(days=30 * i),
                end_date=now + timedelta(days=30 * i + 14),
                is_active=False,
                results_published=False
            )
            db.add(e)
            elections.append(e)

        db.commit()
        for e in elections:
            db.refresh(e)

        # Create 2-3 candidates per election
        all_candidates = {}
        for e in elections:
            cands = []
            for j in range(1, 4):
                c = Candidate(
                    election_id=e.id,
                    name=f"Candidate {j} for {e.name}",
                    description="Candidate from seed",
                    order=j
                )
                db.add(c)
                cands.append(c)
            db.commit()
            for c in cands:
                db.refresh(c)
            all_candidates[e.id] = cands

        # Insert votes: Distribute votes among past/present/future
        # For the CNI '1234' -> 3 past, 1 present, 2 future
        # For '12345' and '123456' do similar but not the exact same elections
        def create_vote(voter_obj, election_obj, candidate_obj, when):
            vote = Vote(
                voter_id=voter_obj.id,
                election_id=election_obj.id,
                candidate_id=candidate_obj.id,
                encrypted_vote=b"seeded_encrypted",
                zero_knowledge_proof=b"seeded_proof",
                ballot_hash=f"seed-{voter_obj.citizen_id}-{election_obj.id}-{candidate_obj.id}",
                timestamp=when
            )
            db.add(vote)

        # Helper to pick candidate
        import random

        # Map citizen to required vote counts
        plan = {
            "1234": {"past": 3, "present": 1, "future": 2},
            "12345": {"past": 2, "present": 1, "future": 2},
            "123456": {"past": 1, "present": 1, "future": 3},
        }

        # Collect elections by status
        past_elections = [e for e in elections if e.end_date < now]
        present_elections = [e for e in elections if e.start_date <= now < e.end_date]
        future_elections = [e for e in elections if e.start_date > now]

        for voter in voters:
            p = plan.get(voter.citizen_id, {"past": 1, "present": 0, "future": 1})

            # Past votes
            for _ in range(p.get("past", 0)):
                if not past_elections:
                    break
                e = random.choice(past_elections)
                c = random.choice(all_candidates[e.id])
                create_vote(voter, e, c, when=e.end_date - timedelta(days=5))

            # Present votes
            for _ in range(p.get("present", 0)):
                if not present_elections:
                    break
                e = random.choice(present_elections)
                c = random.choice(all_candidates[e.id])
                create_vote(voter, e, c, when=now)

            # Future votes (we still create them as scheduled/placeholder votes)
            for _ in range(p.get("future", 0)):
                if not future_elections:
                    break
                e = random.choice(future_elections)
                c = random.choice(all_candidates[e.id])
                create_vote(voter, e, c, when=e.start_date + timedelta(days=1))

        db.commit()

        # Report
        total_voters = db.query(Voter).count()
        total_elections = db.query(Election).count()
        total_votes = db.query(Vote).count()
        print(f"Seeded: voters={total_voters}, elections={total_elections}, votes={total_votes}")

    finally:
        db.close()


def main():
    reset_db()
    seed()


if __name__ == "__main__":
    main()