from flask_wtf import FlaskForm
from wtforms import (StringField, BooleanField, SelectField, TextAreaField)
from wtforms.validators import DataRequired, Length, Optional

from wtforms_sqlalchemy.fields import QuerySelectMultipleField

from common.models.document import Retriever
from common.models.interaction import EveAITool
from common.extensions import cache_manager

from .dynamic_form_base import DynamicFormBase


def get_retrievers():
    return Retriever.query.all()


def get_tools():
    return EveAITool.query.all()


class SpecialistForm(FlaskForm):
    name = StringField('Name', validators=[DataRequired(), Length(max=50)])
    description = TextAreaField('Description', validators=[DataRequired()])

    retrievers = QuerySelectMultipleField(
        'Retrievers',
        query_factory=get_retrievers,
        get_label='name',  # Assuming your Retriever model has a 'name' field
        allow_blank=True,
        description='Select one or more retrievers to associate with this specialist'
    )

    type = SelectField('Specialist Type', validators=[DataRequired()])

    tuning = BooleanField('Enable Specialist Tuning', default=False)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        types_dict = cache_manager.specialist_config_cache.get_types()
        # Dynamically populate the 'type' field using the constructor
        self.type.choices = [(key, value['name']) for key, value in types_dict.items()]


class EditSpecialistForm(DynamicFormBase):
    name = StringField('Name', validators=[DataRequired()])
    description = TextAreaField('Description', validators=[DataRequired()])

    retrievers = QuerySelectMultipleField(
        'Retrievers',
        query_factory=get_retrievers,
        get_label='name',
        allow_blank=True,
        description='Select one or more retrievers to associate with this specialist'
    )

    type = StringField('Specialist Type', validators=[DataRequired()], render_kw={'readonly': True})
    tuning = BooleanField('Enable Retrieval Tuning', default=False)


class BaseComponentForm(DynamicFormBase):
    """Base form for all processing components"""
    name = StringField('Name', validators=[DataRequired(), Length(max=50)])
    description = TextAreaField('Description', validators=[Optional()])
    type = SelectField('Type', validators=[DataRequired()])
    tuning = BooleanField('Enable Tuning', default=False)

    def __init__(self, *args, type_config=None, **kwargs):
        super().__init__(*args, **kwargs)
        if type_config:
            self.type.choices = [(key, value['name']) for key, value in type_config.items()]


# Edit forms that support dynamic fields
class BaseEditComponentForm(DynamicFormBase):
    name = StringField('Name', validators=[DataRequired()])
    description = TextAreaField('Description', validators=[Optional()])
    type = StringField('Type', validators=[DataRequired()], render_kw={'readonly': True})
    tuning = BooleanField('Enable Tuning', default=False)


class EveAIAgentForm(BaseComponentForm):
    role = TextAreaField('Role', validators=[DataRequired()])
    goal = TextAreaField('Goal', validators=[DataRequired()])
    backstory = TextAreaField('Backstory', validators=[DataRequired()])

    tools = QuerySelectMultipleField(
        'Tools',
        query_factory=get_tools,
        get_label='name',
        allow_blank=True,
        description='Select one or more tools that can be used this agent'
    )

    def __init__(self, *args, type_config=None, **kwargs):
        super().__init__(*args, **kwargs)
        if type_config:
            self.type.choices = [(key, value['name']) for key, value in type_config.items()]


class EditEveAIAgentForm(BaseEditComponentForm):
    role = StringField('Role', validators=[DataRequired()])
    goal = StringField('Goal', validators=[DataRequired()])
    backstory = StringField('Backstory', validators=[DataRequired()])

    tools = QuerySelectMultipleField(
        'Tools',
        query_factory=get_tools,
        get_label='name',
        allow_blank=True,
        description='Select one or more tools that can be used this agent'
    )


class EveAITaskForm(BaseComponentForm):
    expected_output = TextAreaField('Expected Output', validators=[DataRequired()])