import math
from typing import Optional, Type, List
from datetime import date, datetime
import secrets
from typing import Type, List
from datetime import date, timezone

from django.db import models
from django.contrib.sessions.backends.db import SessionStore
from pgvector.django import CosineDistance
from django.db.models.manager import BaseManager
from django.db.models import Q
from torch import Tensor

# Import sync_to_async from Django Channels
from asgiref.sync import sync_to_async

from fastapi import HTTPException

from database.models import (
    KhojUser,
    GoogleUser,
    KhojApiUser,
    NotionConfig,
    GithubConfig,
    Entry,
    GithubRepoConfig,
    Conversation,
    ChatModelOptions,
    SearchModelConfig,
    Subscription,
    UserConversationConfig,
    OpenAIProcessorConversationConfig,
    OfflineChatProcessorConversationConfig,
)
from khoj.utils.helpers import generate_random_name
from khoj.search_filter.word_filter import WordFilter
from khoj.search_filter.file_filter import FileFilter
from khoj.search_filter.date_filter import DateFilter


async def set_notion_config(token: str, user: KhojUser):
    notion_config = await NotionConfig.objects.filter(user=user).afirst()
    if not notion_config:
        notion_config = await NotionConfig.objects.acreate(token=token, user=user)
    else:
        notion_config.token = token
        await notion_config.asave()
    return notion_config


async def create_khoj_token(user: KhojUser, name=None):
    "Create Khoj API key for user"
    token = f"kk-{secrets.token_urlsafe(32)}"
    name = name or f"{generate_random_name().title()}"
    return await KhojApiUser.objects.acreate(token=token, user=user, name=name)


def get_khoj_tokens(user: KhojUser):
    "Get all Khoj API keys for user"
    return list(KhojApiUser.objects.filter(user=user))


async def delete_khoj_token(user: KhojUser, token: str):
    "Delete Khoj API Key for user"
    await KhojApiUser.objects.filter(token=token, user=user).adelete()


async def get_or_create_user(token: dict) -> KhojUser:
    user = await get_user_by_token(token)
    if not user:
        user = await create_user_by_google_token(token)
    return user


async def create_user_by_google_token(token: dict) -> KhojUser:
    user, _ = await KhojUser.objects.filter(email=token.get("email")).aupdate_or_create(
        defaults={"username": token.get("email"), "email": token.get("email")}
    )
    await user.asave()

    await GoogleUser.objects.acreate(
        sub=token.get("sub"),
        azp=token.get("azp"),
        email=token.get("email"),
        name=token.get("name"),
        given_name=token.get("given_name"),
        family_name=token.get("family_name"),
        picture=token.get("picture"),
        locale=token.get("locale"),
        user=user,
    )

    await Subscription.objects.acreate(user=user, type="trial")

    return user


def get_user_subscription(email: str) -> Optional[Subscription]:
    return Subscription.objects.filter(user__email=email).first()


async def set_user_subscription(
    email: str, is_recurring=None, renewal_date=None, type="standard"
) -> Optional[Subscription]:
    user_subscription = await Subscription.objects.filter(user__email=email).afirst()
    if not user_subscription:
        user = await get_user_by_email(email)
        if not user:
            return None
        user_subscription = await Subscription.objects.acreate(
            user=user, type=type, is_recurring=is_recurring, renewal_date=renewal_date
        )
        return user_subscription
    elif user_subscription:
        user_subscription.type = type
        if is_recurring is not None:
            user_subscription.is_recurring = is_recurring
        if renewal_date is False:
            user_subscription.renewal_date = None
        elif renewal_date is not None:
            user_subscription.renewal_date = renewal_date
        await user_subscription.asave()
        return user_subscription
    else:
        return None


def get_user_subscription_state(email: str) -> str:
    """Get subscription state of user
    Valid state transitions: trial -> subscribed <-> unsubscribed OR expired
    """
    user_subscription = Subscription.objects.filter(user__email=email).first()
    if not user_subscription:
        return "trial"
    elif user_subscription.type == Subscription.Type.TRIAL:
        return "trial"
    elif user_subscription.is_recurring and user_subscription.renewal_date >= datetime.now(tz=timezone.utc):
        return "subscribed"
    elif not user_subscription.is_recurring and user_subscription.renewal_date >= datetime.now(tz=timezone.utc):
        return "unsubscribed"
    elif not user_subscription.is_recurring and user_subscription.renewal_date < datetime.now(tz=timezone.utc):
        return "expired"
    return "invalid"


async def get_user_by_email(email: str) -> KhojUser:
    return await KhojUser.objects.filter(email=email).afirst()


async def get_user_by_token(token: dict) -> KhojUser:
    google_user = await GoogleUser.objects.filter(sub=token.get("sub")).select_related("user").afirst()
    if not google_user:
        return None
    return google_user.user


async def retrieve_user(session_id: str) -> KhojUser:
    session = SessionStore(session_key=session_id)
    if not await sync_to_async(session.exists)(session_key=session_id):
        raise HTTPException(status_code=401, detail="Invalid session")
    session_data = await sync_to_async(session.load)()
    user = await KhojUser.objects.filter(id=session_data.get("_auth_user_id")).afirst()
    if not user:
        raise HTTPException(status_code=401, detail="Invalid user")
    return user


def get_all_users() -> BaseManager[KhojUser]:
    return KhojUser.objects.all()


def get_user_github_config(user: KhojUser):
    config = GithubConfig.objects.filter(user=user).prefetch_related("githubrepoconfig").first()
    return config


def get_user_notion_config(user: KhojUser):
    config = NotionConfig.objects.filter(user=user).first()
    return config


async def set_text_content_config(user: KhojUser, object: Type[models.Model], updated_config):
    deduped_files = list(set(updated_config.input_files)) if updated_config.input_files else None
    deduped_filters = list(set(updated_config.input_filter)) if updated_config.input_filter else None
    await object.objects.filter(user=user).adelete()
    await object.objects.acreate(
        input_files=deduped_files,
        input_filter=deduped_filters,
        index_heading_entries=updated_config.index_heading_entries,
        user=user,
    )


async def set_user_github_config(user: KhojUser, pat_token: str, repos: list):
    config = await GithubConfig.objects.filter(user=user).afirst()

    if not config:
        config = await GithubConfig.objects.acreate(pat_token=pat_token, user=user)
    else:
        config.pat_token = pat_token
        await config.asave()
        await config.githubrepoconfig.all().adelete()

    for repo in repos:
        await GithubRepoConfig.objects.acreate(
            name=repo["name"], owner=repo["owner"], branch=repo["branch"], github_config=config
        )
    return config


def get_or_create_search_model():
    search_model = SearchModelConfig.objects.filter().first()
    if not search_model:
        search_model = SearchModelConfig.objects.create()

    return search_model


class ConversationAdapters:
    @staticmethod
    def get_conversation_by_user(user: KhojUser):
        conversation = Conversation.objects.filter(user=user)
        if conversation.exists():
            return conversation.first()
        return Conversation.objects.create(user=user)

    @staticmethod
    async def aget_conversation_by_user(user: KhojUser):
        conversation = Conversation.objects.filter(user=user)
        if await conversation.aexists():
            return await conversation.afirst()
        return await Conversation.objects.acreate(user=user)

    @staticmethod
    def has_any_conversation_config(user: KhojUser):
        return ChatModelOptions.objects.filter(user=user).exists()

    @staticmethod
    def get_openai_conversation_config():
        return OpenAIProcessorConversationConfig.objects.filter().first()

    @staticmethod
    async def aget_openai_conversation_config():
        return await OpenAIProcessorConversationConfig.objects.filter().afirst()

    @staticmethod
    def get_offline_chat_conversation_config():
        return OfflineChatProcessorConversationConfig.objects.filter().first()

    @staticmethod
    async def aget_offline_chat_conversation_config():
        return await OfflineChatProcessorConversationConfig.objects.filter().afirst()

    @staticmethod
    def has_valid_offline_conversation_config():
        return OfflineChatProcessorConversationConfig.objects.filter(enabled=True).exists()

    @staticmethod
    def has_valid_openai_conversation_config():
        return OpenAIProcessorConversationConfig.objects.filter().exists()

    @staticmethod
    async def aset_user_conversation_processor(user: KhojUser, conversation_processor_config_id: int):
        config = await ChatModelOptions.objects.filter(id=conversation_processor_config_id).afirst()
        if not config:
            return None
        new_config = await UserConversationConfig.objects.aupdate_or_create(user=user, defaults={"setting": config})
        return new_config

    @staticmethod
    def get_conversation_config(user: KhojUser):
        config = UserConversationConfig.objects.filter(user=user).first()
        if not config:
            return None
        return config.setting

    @staticmethod
    async def aget_conversation_config(user: KhojUser):
        config = await UserConversationConfig.objects.filter(user=user).prefetch_related("setting").afirst()
        if not config:
            return None
        return config.setting

    @staticmethod
    def get_default_conversation_config():
        return ChatModelOptions.objects.filter().first()

    @staticmethod
    async def aget_default_conversation_config():
        return await ChatModelOptions.objects.filter().afirst()

    @staticmethod
    def save_conversation(user: KhojUser, conversation_log: dict):
        conversation = Conversation.objects.filter(user=user)
        if conversation.exists():
            conversation.update(conversation_log=conversation_log)
        else:
            Conversation.objects.create(user=user, conversation_log=conversation_log)

    @staticmethod
    def get_conversation_processor_options():
        return ChatModelOptions.objects.all()

    @staticmethod
    def set_conversation_processor_config(user: KhojUser, new_config: ChatModelOptions):
        user_conversation_config, _ = UserConversationConfig.objects.get_or_create(user=user)
        user_conversation_config.setting = new_config
        user_conversation_config.save()

    @staticmethod
    def has_offline_chat():
        return OfflineChatProcessorConversationConfig.objects.filter(enabled=True).exists()

    @staticmethod
    async def ahas_offline_chat():
        return await OfflineChatProcessorConversationConfig.objects.filter(enabled=True).aexists()

    @staticmethod
    async def get_offline_chat():
        return await ChatModelOptions.objects.filter(model_type="offline").afirst()

    @staticmethod
    async def aget_user_conversation_config(user: KhojUser):
        config = await UserConversationConfig.objects.filter(user=user).prefetch_related("setting").afirst()
        if not config:
            return None
        return config.setting

    @staticmethod
    async def has_openai_chat():
        return await OpenAIProcessorConversationConfig.objects.filter().aexists()

    @staticmethod
    async def get_openai_chat():
        return await ChatModelOptions.objects.filter(model_type="openai").afirst()

    @staticmethod
    async def get_openai_chat_config():
        return await OpenAIProcessorConversationConfig.objects.filter().afirst()


class EntryAdapters:
    word_filer = WordFilter()
    file_filter = FileFilter()
    date_filter = DateFilter()

    @staticmethod
    def does_entry_exist(user: KhojUser, hashed_value: str) -> bool:
        return Entry.objects.filter(user=user, hashed_value=hashed_value).exists()

    @staticmethod
    def delete_entry_by_file(user: KhojUser, file_path: str):
        deleted_count, _ = Entry.objects.filter(user=user, file_path=file_path).delete()
        return deleted_count

    @staticmethod
    def delete_all_entries_by_type(user: KhojUser, file_type: str = None):
        if file_type is None:
            deleted_count, _ = Entry.objects.filter(user=user).delete()
        else:
            deleted_count, _ = Entry.objects.filter(user=user, file_type=file_type).delete()
        return deleted_count

    @staticmethod
    def delete_all_entries(user: KhojUser, file_source: str = None):
        if file_source is None:
            deleted_count, _ = Entry.objects.filter(user=user).delete()
        else:
            deleted_count, _ = Entry.objects.filter(user=user, file_source=file_source).delete()
        return deleted_count

    @staticmethod
    def get_existing_entry_hashes_by_file(user: KhojUser, file_path: str):
        return Entry.objects.filter(user=user, file_path=file_path).values_list("hashed_value", flat=True)

    @staticmethod
    def delete_entry_by_hash(user: KhojUser, hashed_values: List[str]):
        Entry.objects.filter(user=user, hashed_value__in=hashed_values).delete()

    @staticmethod
    def get_entries_by_date_filter(entry: BaseManager[Entry], start_date: date, end_date: date):
        return entry.filter(
            entrydates__date__gte=start_date,
            entrydates__date__lte=end_date,
        )

    @staticmethod
    def user_has_entries(user: KhojUser):
        return Entry.objects.filter(user=user).exists()

    @staticmethod
    async def auser_has_entries(user: KhojUser):
        return await Entry.objects.filter(user=user).aexists()

    @staticmethod
    async def adelete_entry_by_file(user: KhojUser, file_path: str):
        return await Entry.objects.filter(user=user, file_path=file_path).adelete()

    @staticmethod
    def aget_all_filenames_by_source(user: KhojUser, file_source: str):
        return (
            Entry.objects.filter(user=user, file_source=file_source)
            .distinct("file_path")
            .values_list("file_path", flat=True)
        )

    @staticmethod
    async def adelete_all_entries(user: KhojUser):
        return await Entry.objects.filter(user=user).adelete()

    @staticmethod
    def apply_filters(user: KhojUser, query: str, file_type_filter: str = None):
        q_filter_terms = Q()

        explicit_word_terms = EntryAdapters.word_filer.get_filter_terms(query)
        file_filters = EntryAdapters.file_filter.get_filter_terms(query)
        date_filters = EntryAdapters.date_filter.get_query_date_range(query)

        if len(explicit_word_terms) == 0 and len(file_filters) == 0 and len(date_filters) == 0:
            return Entry.objects.filter(user=user)

        for term in explicit_word_terms:
            if term.startswith("+"):
                q_filter_terms &= Q(raw__icontains=term[1:])
            elif term.startswith("-"):
                q_filter_terms &= ~Q(raw__icontains=term[1:])

        q_file_filter_terms = Q()

        if len(file_filters) > 0:
            for term in file_filters:
                q_file_filter_terms |= Q(file_path__regex=term)

            q_filter_terms &= q_file_filter_terms

        if len(date_filters) > 0:
            min_date, max_date = date_filters
            if min_date is not None:
                # Convert the min_date timestamp to yyyy-mm-dd format
                formatted_min_date = date.fromtimestamp(min_date).strftime("%Y-%m-%d")
                q_filter_terms &= Q(embeddings_dates__date__gte=formatted_min_date)
            if max_date is not None:
                # Convert the max_date timestamp to yyyy-mm-dd format
                formatted_max_date = date.fromtimestamp(max_date).strftime("%Y-%m-%d")
                q_filter_terms &= Q(embeddings_dates__date__lte=formatted_max_date)

        relevant_entries = Entry.objects.filter(user=user).filter(
            q_filter_terms,
        )
        if file_type_filter:
            relevant_entries = relevant_entries.filter(file_type=file_type_filter)
        return relevant_entries

    @staticmethod
    def search_with_embeddings(
        user: KhojUser,
        embeddings: Tensor,
        max_results: int = 10,
        file_type_filter: str = None,
        raw_query: str = None,
        max_distance: float = math.inf,
    ):
        relevant_entries = EntryAdapters.apply_filters(user, raw_query, file_type_filter)
        relevant_entries = relevant_entries.filter(user=user).annotate(
            distance=CosineDistance("embeddings", embeddings)
        )
        relevant_entries = relevant_entries.filter(distance__lte=max_distance)

        if file_type_filter:
            relevant_entries = relevant_entries.filter(file_type=file_type_filter)
        relevant_entries = relevant_entries.order_by("distance")
        return relevant_entries[:max_results]

    @staticmethod
    def get_unique_file_types(user: KhojUser):
        return Entry.objects.filter(user=user).values_list("file_type", flat=True).distinct()

    @staticmethod
    def get_unique_file_sources(user: KhojUser):
        return Entry.objects.filter(user=user).values_list("file_source", flat=True).distinct().all()