database/init_db.py

"""
Initialize database - create tables and add sample data
"""
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config.database import engine, Base, SessionLocal
from database.models import *
from datetime import datetime, date, timedelta
import hashlib

def create_tables():
    """Create all database tables"""
    print("Creating database tables...")
    Base.metadata.create_all(bind=engine)
    print("✅ Tables created successfully!")

def add_sample_data():
    """Add sample data for testing"""
    db = SessionLocal()
    
    try:
        # Check if data already exists
        existing_users = db.query(User).first()
        if existing_users:
            print("⚠️ Database already has data. Skipping sample data.")
            return
        
        print("Adding sample data...")
        
        # Create sample therapist user
        therapist_user = User(
            email="therapist@example.com",
            username="dr_smith",
            hashed_password=hashlib.sha256("password123".encode()).hexdigest(),
            first_name="Sarah",
            last_name="Smith",
            role=UserRole.THERAPIST,
            is_verified=True
        )
        db.add(therapist_user)
        db.flush()
        
        # Create therapist profile
        therapist = Therapist(
            user_id=therapist_user.id,
            license_number="SLP12345",
            specialization="Stuttering and Fluency",
            clinic_name="Speech Wellness Center",
            phone="555-0100",
            bio="Specialized in stuttering therapy with 10 years of experience"
        )
        db.add(therapist)
        db.flush()
        
        # Create sample patient users
        patient_user1 = User(
            email="patient1@example.com",
            username="john_doe",
            hashed_password=hashlib.sha256("password123".encode()).hexdigest(),
            first_name="John",
            last_name="Doe",
            role=UserRole.PATIENT,
            is_verified=True
        )
        db.add(patient_user1)
        db.flush()
        
        # Create patient profile
        patient1 = Patient(
            user_id=patient_user1.id,
            therapist_id=therapist.id,
            date_of_birth=date(2000, 5, 15),
            age=24,
            gender="Male",
            diagnosis="Developmental stuttering",
            severity_level="Moderate",
            therapy_goals=["Reduce stuttering frequency", "Improve fluency in conversations"],
            preferred_technique="prolonged_speech"
        )
        db.add(patient1)
        db.flush()
        
        # Create sample assignment
        assignment = Assignment(
            therapist_id=therapist.id,
            patient_id=patient1.id,
            title="Daily Conversation Practice",
            description="Practice conversational speech using prolonged speech technique",
            practice_type="conversation",
            technique="prolonged_speech",
            target_sessions=5,
            target_duration_minutes=10,
            parameters={"topic": "hobbies", "difficulty": "intermediate"},
            due_date=datetime.now() + timedelta(days=7)
        )
        db.add(assignment)
        
        # Create a sample completed session
        session = PracticeSession(
            patient_id=patient1.id,
            assignment_id=assignment.id,
            practice_type="conversation",
            technique="prolonged_speech",
            duration_seconds=300,
            prompt_text="Tell me about your favorite hobby",
            transcribed_text="I really enjoy playing guitar in my free time",
            total_words=10,
            words_per_minute=120,
            is_completed=True,
            completed_at=datetime.now()
        )
        db.add(session)
        db.flush()
        
        # Create analysis for the session
        analysis = SessionAnalysis(
            session_id=session.id,
            stutter_analysis={"total_stutters": 2, "types": {"repetition": 1, "block": 1}},
            fluency_analysis={"prolongation_percentage": 75, "consistency": 0.8},
            total_stutters=2,
            stutter_frequency_percent=5.0,
            fluency_score=75,
            rushed_speech_severity="minimal",
            struggled_phonemes=["g", "t"],
            patient_feedback="Good use of prolonged speech technique. Keep practicing smooth transitions."
        )
        db.add(analysis)
        
        db.commit()
        print("✅ Sample data added successfully!")
        
    except Exception as e:
        print(f"❌ Error adding sample data: {e}")
        db.rollback()
    finally:
        db.close()

def reset_database():
    """Drop all tables and recreate"""
    print("⚠️ Dropping all tables...")
    Base.metadata.drop_all(bind=engine)
    create_tables()
    add_sample_data()

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Database initialization")
    parser.add_argument("--reset", action="store_true", help="Drop and recreate all tables")
    parser.add_argument("--sample-data", action="store_true", help="Add sample data")
    
    args = parser.parse_args()
    
    if args.reset:
        response = input("⚠️ This will DELETE all data. Are you sure? (yes/no): ")
        if response.lower() == "yes":
            reset_database()
        else:
            print("Cancelled.")
    else:
        create_tables()
        if args.sample_data:
            add_sample_data()