from telethon import TelegramClient as BaseTelegramClient
from .session import SessionManager
from .parsing.manager import EntitiesManager


class TelegramClient(BaseTelegramClient):
    def __init__(
        self,
        session,
        api_id,
        api_hash,
        *,
        device_model=None,
        system_version=None,
        app_version=None,
        lang_code="en",
        system_lang_code="en",
        **kwargs,
    ):
        super().__init__(session, api_id, api_hash, **kwargs)
        self._session_manager = SessionManager(self.session)
        self._entities_manager = EntitiesManager()

        self._device_model = device_model
        self._system_version = system_version
        self._app_version = app_version
        self._lang_code = lang_code
        self._system_lang_code = system_lang_code

    async def connect(self):
        self._session_manager.init_session(
            device_model=self._device_model,
            system_version=self._system_version,
            app_version=self._app_version,
            lang_code=self._lang_code,
            system_lang_code=self._system_lang_code,
        )
        await super().connect()

    async def _process_entities(self, text, formatting_entities, parse_mode):
        if not text:
            return text, formatting_entities, parse_mode

        text, new_entities = await self._entities_manager.parse(text)
        if new_entities:
            formatting_entities = list(formatting_entities or []) + new_entities
            formatting_entities.sort(key=lambda e: e.offset)
            parse_mode = None

        return text, formatting_entities, parse_mode

    async def send_message(
        self,
        entity,
        message,
        *,
        parse_mode=(),
        formatting_entities=None,
        **kwargs,
    ):
        message, formatting_entities, parse_mode = await self._process_entities(
            message, formatting_entities, parse_mode
        )
        return await super().send_message(
            entity,
            message,
            parse_mode=parse_mode,
            formatting_entities=formatting_entities,
            **kwargs,
        )

    async def edit_message(
        self,
        entity,
        message=None,
        *,
        parse_mode=(),
        formatting_entities=None,
        **kwargs,
    ):
        message, formatting_entities, parse_mode = await self._process_entities(
            message, formatting_entities, parse_mode
        )
        return await super().edit_message(
            entity,
            message,
            parse_mode=parse_mode,
            formatting_entities=formatting_entities,
            **kwargs,
        )

    async def send_file(
        self,
        entity,
        file,
        *,
        caption=None,
        parse_mode=(),
        formatting_entities=None,
        **kwargs,
    ):
        caption, formatting_entities, parse_mode = await self._process_entities(
            caption, formatting_entities, parse_mode
        )
        return await super().send_file(
            entity,
            file,
            caption=caption,
            parse_mode=parse_mode,
            formatting_entities=formatting_entities,
            **kwargs,
        )

    async def parse_telecats(self, text: str, remove_only=False):
        if remove_only:
            return await self._entities_manager.remove(text)
        return await self._entities_manager.parse(text)
