forms.py

from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField, SelectField, TextAreaField, DateField
from wtforms.validators import DataRequired, Length, EqualTo, ValidationError, Optional
from .models import User # User is already imported, good.

class RegistrationForm(FlaskForm):
    username = StringField('Username', 
                           validators=[DataRequired(), Length(min=2, max=80)]) # Max length matches User model
    password = PasswordField('Password', 
                             validators=[DataRequired(), Length(min=6)])
    confirm_password = PasswordField('Confirm Password', 
                                     validators=[DataRequired(), EqualTo('password')])
    role = SelectField('Role', 
                       choices=[('employee', 'Employee'), ('manager', 'Manager')],
                       validators=[DataRequired()])
    submit = SubmitField('Sign Up')

    def validate_username(self, username):
        user = User.query.filter_by(username=username.data).first()
        if user:
            raise ValidationError('That username is already taken. Please choose a different one.')

class LoginForm(FlaskForm):
    username = StringField('Username', 
                           validators=[DataRequired(), Length(min=2, max=80)]) # Max length matches User model
    password = PasswordField('Password', validators=[DataRequired()])
    submit = SubmitField('Login')

class TaskForm(FlaskForm):
    title = StringField('Title', validators=[DataRequired()])
    description = TextAreaField('Description')
    due_date = DateField('Due Date', format='%Y-%m-%d', validators=[Optional()])
    status = SelectField('Status', 
                         choices=[('pending', 'Pending'), 
                                  ('in progress', 'In Progress'), 
                                  ('completed', 'Completed')], 
                         validators=[DataRequired()]) # Default is typically the first choice
    assignee_id = SelectField('Assign To', coerce=int, validators=[Optional()]) # Added Optional for now, will adjust in __init__
    submit = SubmitField('Save Task')

    def __init__(self, *args, **kwargs):
        current_user_role = kwargs.pop('current_user_role', None)
        super(TaskForm, self).__init__(*args, **kwargs)

        if current_user_role == 'manager':
            # Populate choices for assignee_id
            employees = User.query.filter_by(role='employee').all()
            self.assignee_id.choices = [(user.id, user.username) for user in employees]
            # Add an option for 'Unassigned' or 'Assign to Self (Manager)' if desired
            # For now, let's make it required to select an employee if manager is creating/editing
            self.assignee_id.choices.insert(0, (0, 'Unassigned / Assign to Self')) # Representing unassigned or manager self-assignment
            self.assignee_id.validators = [DataRequired()] # Make it required for manager
        else:
            # For employees, this field might be hidden or disabled in the template.
            # Or, we can remove it from the form items if it's not relevant.
            # For now, let's give it a default choice that makes sense if it were to be submitted.
            self.assignee_id.choices = [] # No choices for employee to change assignee
            # If we want to ensure it's not submitted by employees, we can clear validators or set a default.
            # However, the route logic will handle setting assignee_id to current_user.id for employees.
            # So, an empty choices list and Optional validator is fine.
            # If current_user is available here, we could set a default choice:
            # self.assignee_id.choices = [(current_user.id, current_user.username)]
            # self.assignee_id.data = current_user.id
            pass # Keep validators as Optional for non-managers