import json

from sqlalchemy import (
    Column,
    Integer,
    String,
)
from consts.consts import GPT_DEFAULT, LANGUAGE_EN, Claude3, EmbeddingModel, KnowledgeSourceType, SuggestQuestionLevel
from utils.local import get_tl_override_ai_model
from utils.logging import logging
from models.base import Base, BaseMixin, with_session


def is_internal_api_enabled(email: str, org_id=""):
    preference_by_email_suffix = UserPreference.load_by_email_suffix(email.split('@')[-1])
    preference_by_org_id = UserPreference.load_by_org_id(org_id)

    return preference_by_email_suffix.internal_api_enabled or preference_by_org_id.internal_api_enabled


class UserPreference:
    def __init__(self) -> None:
        self.llm_api_key = None
        self.language_code = LANGUAGE_EN
        self.llm_model = Claude3.DEFAULT
        self.mode_type = ""

        self.embedding_model = EmbeddingModel.ADA_002
        self.kg_source_type = KnowledgeSourceType.LOCAL
        self.assistant_enabled = False

        self.rate_limit_per_day = 100
        self.internal_api_enabled = False

        # suggest question
        self.suggest_question_level = SuggestQuestionLevel.SIMPLE

    def __str__(self) -> str:
        return f'UserPreference: {self.__dict__}'

    @classmethod
    def load(cls, app_id: int):
        item = cls()

        with with_session() as s:
            preference = Preference.get_by_app_id(s, app_id)
            if not preference:
                logging.info('app_id = %s, load user preference, setted by default, preferences = %s', app_id, str(item.to_json()))
                return item

            for k, v in preference.items():
                setattr(item, k, v)

            logging.info('app_id = %s, load user preference, setted by user, preferences = %s', app_id, str(item.to_json()))

            override_ai_model = get_tl_override_ai_model()
            if override_ai_model:
                logging.info('app_id = %s, load user preference, override ai model = %s', app_id, override_ai_model)
                item.llm_model = override_ai_model

            return item

    @classmethod
    def load_by_org_id(cls, org_id: str):
        item = cls()

        with with_session() as s:
            preference = Preference.get_by_org_id(s, org_id)
            if not preference:
                logging.info('org_id = %s, load user preference, setted by default, preferences = %s', org_id, str(item.to_json()))
                return item

            for k, v in preference.items():
                setattr(item, k, v)

            logging.info('org_id = %s, load user preference, setted by user, preferences = %s', org_id, str(item.to_json()))
            return item

    @classmethod
    def load_by_email_suffix(cls, email_suffix: str):
        item = cls()

        with with_session() as s:
            preference = Preference.get_by_email_suffix(s, email_suffix)
            if not preference:
                logging.info('email_suffix = %s, load user preference, setted by default, preferences = %s', email_suffix, str(item.to_json()))
                return item

            for k, v in preference.items():
                setattr(item, k, v)

            logging.info('email_suffix = %s, load user preference, setted by user, preferences = %s', email_suffix, str(item.to_json()))
            return item

    def to_json(self):
        return self.__dict__

    def to_history_log(self):
        llm_api_key = self.llm_api_key
        if llm_api_key:
            llm_api_key = llm_api_key[:4] + '*' * 10 + llm_api_key[-4:]

        return {
            "llm_api_key": llm_api_key,
            "language_code": self.language_code,
            "llm_model": self.llm_model,
        }


class Preference(Base, BaseMixin):
    __tablename__ = 'preference'

    org_id = Column(String(64), nullable=False, index=True, default="")
    email_suffix = Column(String(64), nullable=False, index=True, default="")
    app_id = Column(Integer, nullable=False, index=True, default=0)
    preference = Column(String(4096), nullable=False)

    @classmethod
    def upsert(cls, s, app_id: int, preference: dict):
        logging.info("app_id = %s, upsert user preference, preference = %s", app_id, str(preference))

        obj = s.query(cls).filter(cls.app_id == app_id).first()
        if obj:
            value = json.loads(obj.preference)
            value.update(preference)
            obj.preference = json.dumps(value)
        else:
            obj = cls(app_id=app_id, preference=json.dumps(preference))
            s.add(obj)

        s.flush()

        return obj

    @classmethod
    def upsert_by_org_id(cls, s, org_id: str, preference: dict):
        logging.info("org_id = %s, upsert user preference, preference = %s", org_id, str(preference))

        obj = s.query(cls).filter(cls.org_id == org_id).first()
        if obj:
            value = json.loads(obj.preference)
            value.update(preference)
            obj.preference = json.dumps(value)
        else:
            obj = cls(org_id=org_id, preference=json.dumps(preference))
            s.add(obj)

        s.flush()

        return obj

    @classmethod
    def upsert_by_email_suffix(cls, s, email_suffix: str, preference: dict):
        logging.info("email_suffix = %s, upsert user preference, preference = %s", email_suffix, str(preference))

        obj = s.query(cls).filter(cls.email_suffix == email_suffix).first()
        if obj:
            value = json.loads(obj.preference)
            value.update(preference)
            obj.preference = json.dumps(value)
        else:
            obj = cls(email_suffix=email_suffix, preference=json.dumps(preference))
            s.add(obj)

        s.flush()

        return obj

    @classmethod
    def get_by_app_id(cls, s, app_id: int):
        obj = s.query(cls).filter(cls.app_id == app_id).first()
        if obj:
            return json.loads(obj.preference)
        else:
            return {}

    @classmethod
    def delete_preference(cls, s, preference: list, app_id: int = 0, org_id: str = "", email_suffix: str = ""):
        filters = []
        if app_id:
            filters.append(cls.app_id == app_id)
        if org_id:
            filters.append(cls.org_id == org_id)
        if email_suffix:
            filters.append(cls.email_suffix == email_suffix)

        obj = s.query(cls).filter(*filters).first()
        if obj:
            value = json.loads(obj.preference)
            for k in preference:
                if k in value:
                    value.pop(k)
            obj.preference = json.dumps(value)
            s.flush()

    @classmethod
    def get_by_org_id(cls, s, org_id: str):
        obj = s.query(cls).filter(cls.org_id == org_id).first()
        if obj:
            return json.loads(obj.preference)
        else:
            return {}

    @classmethod
    def get_by_email_suffix(cls, s, email_suffix: str):
        obj = s.query(cls).filter(cls.email_suffix == email_suffix).first()
        if obj:
            return json.loads(obj.preference)
        else:
            return {}