# Telebot imports
from telebot.async_telebot import AsyncTeleBot
from telebot.asyncio_storage import StateMemoryStorage, StatePickleStorage
from telebot.asyncio_handler_backends import State, StatesGroup
from telebot.asyncio_filters import StateFilter
from telebot import types
from telebot.callback_data import CallbackData, CallbackDataFilter
from telebot.types import User as TgUser
from telebot.types import Chat as TgChat
from telebot.types import Message
from telebot.util import user_link

# Async things imports
import asyncio

# Other modules imports
from datetime import datetime
import math
from typing import Union

# Local imports
from config import token
import textbook
import keyboards

# DB
from db.base import Session, engine, Base
from db.models import User, Group, GroupMember, Fund, FundMember
from sqlalchemy.orm import joinedload


bot = AsyncTeleBot(token, state_storage=StatePickleStorage())
session: Session = None


class States(StatesGroup):
    default = State()
    newfund_amount = State()
    newfund_description = State()
    close_fund = State()


# Utils


def get_fund_text(fund: Fund):
    count = fund.members.count()
    contributors = fund.members.filter(FundMember.contributed).count()
    personal_amount = math.ceil(fund.amount / count)
    return textbook.fund.format(
        active="🟢" if fund.active else "🔴",
        name=fund.name,
        amount=fund.amount,
        description=fund.description,
        personal_amount=personal_amount,
        collected_amount=personal_amount * contributors,
        contributors=contributors,
        count=count,
    )


def get_user(tg_user: TgUser) -> User:
    if user := session.query(User).filter(User.id == tg_user.id).first():
        return user
    user = User(
        id=tg_user.id,
        name=tg_user.first_name,
        username=tg_user.username,
    )
    session.add(user)
    session.commit()
    return user


def get_group(tg_chat: TgChat) -> Group:
    if group := session.query(Group).filter(Group.id == tg_chat.id).first():
        return group
    group = Group(
        id=tg_chat.id,
    )
    session.add(group)
    session.commit()
    return group


# Bot logic


@bot.message_handler(commands=["start"])
async def start(msg: Message):
    if msg.chat.type == "private":
        user = get_user(msg.from_user)
        fund_members_count = (
            session.query(FundMember)
            .join(Fund)
            .filter(Fund.owner_id == user.id, Fund.active == True)
            .options(joinedload(FundMember.fund))
            .distinct(FundMember.fund)
            .count()
        )
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.start_private.format(count=fund_members_count),
        )

    elif msg.chat.type in ("group", "supergroup"):
        get_group(msg.chat)
        await bot.send_message(chat_id=msg.chat.id, text=textbook.start_group)


# @bot.message_handler(commands=["setup_list"])
# async def setup_list(msg: Message):
#     if msg.chat.type in ("group", "supergroup"):
#         group = get_group(msg.chat)
#         members = group.group_members
#         await bot.send_message(
#             chat_id=msg.chat.id,
#             text=textbook.setup_list.format(
#                 count=members.count(),
#                 members=", ".join(m.user.name for m in members.all()),
#             ),
#             parse_mode="html",
#         )


def get_setup_text(group):
    members = group.group_members
    return textbook.setup + textbook.setup_list.format(count=members.count(), members=", ".join(f'<a href="tg://user?id={m.user.id}">{m.user.name}</a>'  for m in members.all()))


@bot.message_handler(commands=["setup"])
async def setup(msg: Message):
    group = get_group(msg.chat)
    if msg.chat.type in ("group", "supergroup"):
        members = group.group_members
        await bot.send_message(
            chat_id=msg.chat.id, 
            text=get_setup_text(group),
            reply_markup=keyboards.setup(),
            parse_mode="html",
        )


@bot.callback_query_handler(func=lambda c: c.data == "register_group_member")
async def register_group_member(call: types.CallbackQuery):
    new = False
    get_user(call.from_user)
    group = get_group(call.message.chat)
    if (
        group_member := session.query(GroupMember)
        .filter(
            GroupMember.group_id == call.message.chat.id,
            GroupMember.user_id == call.from_user.id,
        )
        .first()
    ):
        session.delete(group_member)
    else:
        group_member = GroupMember(
            group_id=call.message.chat.id,
            user_id=call.from_user.id,
        )
        session.add(group_member)
        new = True
    session.commit()
    await bot.edit_message_text(
        chat_id=call.message.chat.id, 
        text=get_setup_text(group),
        reply_markup=keyboards.setup(),
        message_id=call.message.id,
        parse_mode="html",
    )
    await bot.answer_callback_query(
        callback_query_id=call.id,
        text=textbook.user_parted if new else textbook.user_left,
        show_alert=True,
    )


@bot.message_handler(commands=["newfund"])
async def newfund(msg: Message):
    if msg.chat.type in ("group", "supergroup"):
        if group := get_group(msg.chat):
            if group.group_members.count():
                if (
                    session.query(Fund)
                    .filter(Fund.group_id == msg.chat.id, Fund.active)
                    .first()
                ):
                    await bot.send_message(
                        chat_id=msg.chat.id,
                        text=textbook.newfund_already_exists,
                    )
                else:
                    await bot.set_state(
                        user_id=msg.from_user.id,
                        chat_id=msg.chat.id,
                        state=States.newfund_amount,
                    )
                    await bot.send_message(
                        chat_id=msg.chat.id,
                        text=textbook.newfund_amount.format(
                            user=user_link(msg.from_user)
                        ),
                        reply_markup=keyboards.cancel(),
                        parse_mode="html",
                    )
            else:
                await bot.send_message(
                    chat_id=msg.chat.id,
                    text=textbook.not_set_up,
                )


@bot.callback_query_handler(
    func=lambda c: c.data == "cancel", state=States.newfund_amount
)
async def cancel_newfund_amount(call: types.CallbackQuery):
    await bot.set_state(
        user_id=call.from_user.id, chat_id=call.message.chat.id, state=States.default
    )
    await bot.delete_message(chat_id=call.message.chat.id, message_id=call.message.id)
    await bot.answer_callback_query(
        callback_query_id=call.id,
        text=textbook.cancel,
    )


@bot.callback_query_handler(
    func=lambda c: c.data == "cancel", state=States.newfund_description
)
async def cancel_newfund_description(call: types.CallbackQuery):
    await bot.set_state(
        user_id=call.from_user.id, chat_id=call.message.chat.id, state=States.default
    )
    await bot.delete_message(chat_id=call.message.chat.id, message_id=call.message.id)
    await bot.answer_callback_query(
        callback_query_id=call.id,
        text=textbook.cancel,
    )


@bot.message_handler(content_types=["text"], state=States.newfund_amount)
async def newfund_amount(msg: Message):
    if msg.text.isdigit():
        async with bot.retrieve_data(
            user_id=msg.from_user.id, chat_id=msg.chat.id
        ) as state_data:
            state_data["newfund_amount"] = int(msg.text)
        await bot.set_state(
            user_id=msg.from_user.id,
            chat_id=msg.chat.id,
            state=States.newfund_description,
        )
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.newfund_description,
            reply_markup=keyboards.cancel(),
            parse_mode="html",
        )

    else:
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.not_number.format(user=user_link(msg.from_user)),
            reply_markup=keyboards.cancel(),
            parse_mode="html",
        )


@bot.message_handler(content_types=["text"], state=States.newfund_description)
async def newfund_description(msg: Message):
    if len(msg.text) < 121:
        await bot.set_state(
            user_id=msg.from_user.id, chat_id=msg.chat.id, state=States.default
        )
        async with bot.retrieve_data(
            user_id=msg.from_user.id, chat_id=msg.chat.id
        ) as state_data:
            amount = state_data.get("newfund_amount")
        fund = Fund(
            owner_id=msg.from_user.id,
            group_id=msg.chat.id,
            description=msg.text,
            amount=amount,
            name=datetime.today().strftime("%d.%m"),
        )
        session.add(fund)
        for group_member in session.query(GroupMember).filter(
            GroupMember.group_id == msg.chat.id
        ):
            fund_member = FundMember(user_id=group_member.user.id, fund_id=fund.id)
            session.add(fund_member)
        session.commit()
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.fund_created.format(fund=fund.name),
        )
        await asyncio.sleep(1)
        await bot.send_message(
            chat_id=msg.chat.id,
            text=get_fund_text(fund),
            reply_markup=keyboards.fund_markup(),
            parse_mode="html",
        )

    else:
        await bot.send_message(
            chat_id=msg.chat.id,
            text=textbook.newfund_description_too_long,
            reply_markup=keyboards.cancel(),
            parse_mode="html",
        )


@bot.message_handler(commands=["fund"])
async def fund(msg: Message):
    if msg.chat.type in ("group", "supergroup"):
        group = get_group(msg.chat)
        if (
            fund := session.query(Fund)
            .filter(Fund.group_id == group.id, Fund.active == True)
            .first()
        ):
            await bot.send_message(
                chat_id=msg.chat.id,
                text=get_fund_text(fund),
                reply_markup=keyboards.fund_markup(),
                parse_mode="html",
            )
        else:
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.fund_not_found,
            )


@bot.callback_query_handler(func=lambda c: c.data == "close_fund")
async def close_fund_prompt(call: types.CallbackQuery):
    if group := get_group(call.message.chat):
        if fund := group.funds.filter(Fund.active).first():
            if call.from_user.id == fund.owner_id:
                await bot.edit_message_text(
                    text=textbook.close_fund_prompt,
                    chat_id=call.message.chat.id,
                    message_id=call.message.id,
                    reply_markup=keyboards.yes_no(),
                    parse_mode="html",
                )
                await bot.set_state(
                    user_id=call.from_user.id,
                    chat_id=call.message.chat.id,
                    state=States.close_fund,
                )
            else:
                await bot.answer_callback_query(
                    callback_query_id=call.id,
                    text=textbook.not_owner.format(owner_name=fund.owner.name),
                )
        else:
            await bot.send_message(
                chat_id=call.message.chat.id,
                text=textbook.fund_not_found,
            )


@bot.callback_query_handler(
    func=lambda c: c.data in ("yes", "no"), state=States.close_fund
)
async def close_fund(call: types.CallbackQuery):
    if group := get_group(call.message.chat):
        if fund := group.funds.filter(Fund.active).first():
            if call.from_user.id == fund.owner_id:
                if call.data == "yes":
                    setattr(fund, "active", False)
                    session.commit()
                    await bot.edit_message_text(
                        text=textbook.fund_closed.format(
                            name=fund.name, fund_text=get_fund_text(fund)
                        ),
                        chat_id=call.message.chat.id,
                        message_id=call.message.id,
                        parse_mode="html",
                    )
                elif call.data == "no":
                    await bot.edit_message_text(
                        text=get_fund_text(fund),
                        chat_id=call.message.chat.id,
                        message_id=call.message.id,
                        reply_markup=keyboards.fund_markup(),
                        parse_mode="html",
                    )
            else:
                await bot.answer_callback_query(
                    callback_query_id=call.id,
                    text=textbook.not_owner.format(owner_name=fund.owner.name),
                )
        else:
            await bot.send_message(
                chat_id=call.message.chat.id,
                text=textbook.fund_not_found,
            )


@bot.callback_query_handler(func=lambda c: c.data == "contributed")
async def contributed(call: types.CallbackQuery):
    group = get_group(call.message.chat)
    user = get_user(call.from_user)
    if fund := group.funds.filter(Fund.active).first():
        if fund_user := fund.members.filter(
            FundMember.user_id == call.from_user.id
        ).first():
            if not fund_user.contributed:
                setattr(fund_user, "contributed", True)
                session.commit()
                await bot.answer_callback_query(
                    callback_query_id=call.id,
                    text=textbook.contributed,
                )
                await bot.edit_message_text(
                    text=get_fund_text(fund),
                    chat_id=call.message.chat.id,
                    message_id=call.message.id,
                    reply_markup=keyboards.fund_markup(),
                    parse_mode="html",
                )
                if (
                    fund.members.count()
                    == fund.members.filter(FundMember.contributed == True).count()
                ):
                    setattr(fund, "active", False)
                    session.commit()
                    await bot.edit_message_text(
                        text=get_fund_text(fund),
                        chat_id=call.message.chat.id,
                        message_id=call.message.id,
                        reply_markup=None,
                        parse_mode="html",
                    )
                    await bot.send_message(
                        chat_id=call.message.chat.id,
                        text=textbook.fund_completed.format(
                            fund_name=fund.name,
                            owner_str=f'<a href="tg://user?id={fund.owner.id}">{fund.owner.name}</a>',
                        ),
                        parse_mode="html",
                    )
            else:
                await bot.answer_callback_query(
                    callback_query_id=call.id,
                    text=textbook.already_contributed,
                    show_alert=True,
                )

        else:
            await bot.answer_callback_query(
                callback_query_id=call.id,
                text=textbook.not_fund_member,
                show_alert=True,
            )
    else:
        await bot.answer_callback_query(
            callback_query_id=call.id, text=textbook.fund_not_found, show_alert=True
        )


@bot.message_handler(commands=["remind"])
async def remind(msg: Message):
    if msg.chat.type in ("group", "supergroup"):
        group = get_group(msg.chat)
        if (
            fund := session.query(Fund)
            .filter(Fund.group_id == group.id, Fund.active == True)
            .first()
        ):
            not_contributed = fund.members.filter(FundMember.contributed == False).all()
            if len(not_contributed):
                s = ""
                for fm in not_contributed:
                    s += f'<a href="tg://user?id={fm.user.id}">{fm.user.name}</a>\n'
                await bot.send_message(
                    chat_id=msg.chat.id,
                    text=textbook.remind.format(fund_name=fund.name, s=s),
                    parse_mode="html",
                )
            else:
                await bot.send_message(
                    chat_id=msg.chat.id,
                    text=textbook.remind_already.format(fund_name=fund.name),
                    parse_mode="html",
                )
        else:
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.fund_not_found,
            )


@bot.message_handler(commands=["dmremind"])
async def dmremind(msg: Message):
    if msg.chat.type in ("group", "supergroup"):
        group = get_group(msg.chat)
        if (
            fund := session.query(Fund)
            .filter(Fund.group_id == group.id, Fund.active == True)
            .first()
        ):
            not_contributed = fund.members.filter(FundMember.contributed == False).all()
            if len(not_contributed):
                counter = 0
                not_sent = []
                for member in not_contributed:
                    try:
                        await asyncio.sleep(0.1)
                        await bot.send_message(
                            chat_id=member.user.id,
                            text=textbook.dmremind.format(
                                fund_name=fund.name, chat_name=msg.chat.title
                            ),
                        )
                        counter += 1
                    except Exception:
                        not_sent.append(member)
                s = textbook.dmremind_completed.format(
                    sent_count=counter, members_count=len(not_contributed)
                )

                if not_sent:
                    s += "\n" + textbook.dmremind_not_sent_list.format(
                        members=", ".join([m.user.name for m in not_sent])
                    )
                await bot.send_message(
                    chat_id=msg.chat.id,
                    text=s,
                    parse_mode="html",
                )
            else:
                await bot.send_message(
                    chat_id=msg.chat.id,
                    text=textbook.remind_already.format(fund_name=fund.name),
                    parse_mode="html",
                )
        else:
            await bot.send_message(
                chat_id=msg.chat.id,
                text=textbook.fund_not_found,
            )


@bot.message_handler(commands=["mystate"])
async def mystate(msg: Message):
    state = await bot.get_state(user_id=msg.from_user.id)
    await bot.send_message(chat_id=msg.chat.id, text=state)


@bot.message_handler(commands=["chatid"])
async def chatid(msg: Message):
    await bot.send_message(chat_id=msg.chat.id, text=msg.chat.id)


@bot.message_handler(commands=["guide"])
async def guide(msg: Message):
    await bot.send_message(chat_id=msg.chat.id, text=textbook.guide, parse_mode="html")


@bot.message_handler(commands=["commands"])
async def commands(msg: Message):
    await bot.send_message(
        chat_id=msg.chat.id, text=textbook.commands, parse_mode="html"
    )


async def main():
    a = asyncio.create_task(bot.polling(non_stop=True))
    await a


if __name__ == "__main__":
    print("Bot started", flush=True)
    Base.metadata.create_all(engine)
    session = Session()
    bot.add_custom_filter(StateFilter(bot))
    bot.enable_saving_states(filename="./.state-save/states.pkl")
    asyncio.run(main())