from flask import session
from flask_wtf import FlaskForm
from wtforms import (StringField, BooleanField, SelectField, TextAreaField)
from wtforms.fields.datetime import DateField
from wtforms.fields.numeric import IntegerField, FloatField
from wtforms.validators import DataRequired, Length, Optional, NumberRange

from wtforms_sqlalchemy.fields import QuerySelectMultipleField

from common.models.document import Retriever
from common.models.interaction import EveAITool, Specialist
from common.models.user import TenantMake
from common.extensions import cache_manager
from common.utils.form_assistants import validate_json

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=[Optional()])

    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()])
    active = BooleanField('Active', validators=[Optional()], default=True)
    tuning = BooleanField('Enable Specialist Tuning', default=False)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        types_dict = cache_manager.specialists_types_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=[Optional()])
    active = BooleanField('Active', validators=[Optional()], default=True)

    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})
    type_version = StringField('Type Version', validators=[DataRequired()], render_kw={'readonly': True})
    tuning = BooleanField('Enable Specialist 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})
    type_version = StringField('Type Version', validators=[DataRequired()], render_kw={'readonly': True})
    tuning = BooleanField('Enable Tuning', default=False)


class EditEveAIAgentForm(BaseEditComponentForm):
    role = TextAreaField('Role', validators=[Optional()])
    goal = TextAreaField('Goal', validators=[Optional()])
    backstory = TextAreaField('Backstory', validators=[Optional()])
    temperature = FloatField('Temperature', validators=[Optional(), NumberRange(min=0, max=1)])
    llm_model = SelectField('LLM Model', validators=[Optional()])

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        agent_config = cache_manager.agents_config_cache.get_config(self.type, self.type_version)
        if agent_config.get('allowed_models', None):
            self.llm_model.choices = agent_config.allowed_models
        else:
            self.llm_model.choices = agent_config.get('full_model_name', 'mistral.mistral-medium-latest')


class EditEveAITaskForm(BaseEditComponentForm):
    task_description = StringField('Task Description', validators=[Optional()])
    expected_outcome = StringField('Expected Outcome', validators=[Optional()])


class EditEveAIToolForm(BaseEditComponentForm):
    pass


class ExecuteSpecialistForm(DynamicFormBase):
    id = IntegerField('Specialist ID', validators=[DataRequired()], render_kw={'readonly': True})
    name = StringField('Specialist Name', validators=[DataRequired()], render_kw={'readonly': True})
    description = TextAreaField('Specialist Description', validators=[Optional()], render_kw={'readonly': True})


class SpecialistMagicLinkForm(FlaskForm):
    name = StringField('Name', validators=[DataRequired(), Length(max=50)])
    description = TextAreaField('Description', validators=[Optional()])
    magic_link_code = StringField('Magic Link Code', validators=[DataRequired(), Length(max=55)], render_kw={'readonly': True})
    specialist_id = SelectField('Specialist', validators=[DataRequired()])
    valid_from = DateField('Valid From', id='form-control datepicker', validators=[Optional()])
    valid_to = DateField('Valid To', id='form-control datepicker', validators=[Optional()])

    # Metadata fields
    user_metadata = TextAreaField('User Metadata', validators=[Optional(), validate_json])
    system_metadata = TextAreaField('System Metadata', validators=[Optional(), validate_json])

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        specialists = Specialist.query.all()
        # Dynamically populate the specialist field
        self.specialist_id.choices = [(specialist.id, specialist.name) for specialist in specialists]


class EditSpecialistMagicLinkForm(DynamicFormBase):
    name = StringField('Name', validators=[DataRequired(), Length(max=50)])
    description = TextAreaField('Description', validators=[Optional()])
    magic_link_code = StringField('Magic Link Code', validators=[DataRequired(), Length(max=55)],
                                  render_kw={'readonly': True})
    specialist_id = IntegerField('Specialist', validators=[DataRequired()], render_kw={'readonly': True})
    specialist_name = StringField('Specialist Name', validators=[DataRequired()], render_kw={'readonly': True})

    tenant_make_id = SelectField('Tenant Make', validators=[Optional()], coerce=int)
    valid_from = DateField('Valid From', id='form-control datepicker', validators=[Optional()])
    valid_to = DateField('Valid To', id='form-control datepicker', validators=[Optional()])

    # Metadata fields
    user_metadata = TextAreaField('User Metadata', validators=[Optional(), validate_json])
    system_metadata = TextAreaField('System Metadata', validators=[Optional(), validate_json])

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        specialist = Specialist.query.get(kwargs['specialist_id'])
        if specialist:
            self.specialist_name.data = specialist.name
        else:
            self.specialist_name.data = ''

        # Dynamically populate the tenant_make field with None as first option
        tenant_id = session.get('tenant').get('id')
        tenant_makes = TenantMake.query.filter_by(tenant_id=tenant_id).all()
        self.tenant_make_id.choices = [(0, 'None')] + [(make.id, make.name) for make in tenant_makes]


class ViewSpecialistMagicLinkURLsForm(FlaskForm):
    name = StringField('Name', validators=[DataRequired(), Length(max=50)])
    description = TextAreaField('Description', validators=[Optional()])
    magic_link_code = StringField('Magic Link Code', validators=[DataRequired(), Length(max=55)], render_kw={'readonly': True})

    chat_client_url = StringField('Chat Client URL', validators=[Optional()], render_kw={'readonly': True})
    qr_code_url = StringField('QR Code', validators=[Optional()], render_kw={'readonly': True})