from flask_wtf import FlaskForm
from wtforms import StringField, FloatField, SubmitField, SelectField
from wtforms.validators import DataRequired
from app.models import Account

class AccountForm(FlaskForm):
    name = StringField('Account Name', validators=[DataRequired()])
    account_type = SelectField('Account Type', choices=[
        ('asset', 'Asset'),
        ('expense', 'Expense'),
        ('income', 'Income')
    ], validators=[DataRequired()])
    balance = FloatField('Initial Balance', validators=[DataRequired()])
    submit = SubmitField('Add Account')

class TransactionForm(FlaskForm):
    amount = FloatField('Amount', validators=[DataRequired()])
    type = SelectField('Type', choices=[
        ('transfer', 'Transfer'),
        ('expense', 'Expense'),
        ('income', 'Income')
    ], validators=[DataRequired()])
    source_account = SelectField('Source Account', coerce=int, validators=[DataRequired()])
    destination_account = SelectField('Destination Account', coerce=int, validators=[DataRequired()])
    concept = StringField('Concept', validators=[DataRequired()])
    submit = SubmitField('Add Transaction')

    def __init__(self, *args, **kwargs):
        super(TransactionForm, self).__init__(*args, **kwargs)
        self.source_account.choices = [(account.id, account.name) for account in Account.query.all()]
        self.destination_account.choices = [(account.id, account.name) for account in Account.query.all()]