import os
import asyncio
from typing import Optional, Callable, Any
from telethon import TelegramClient
from telethon.sessions import StringSession
from telethon.errors import SessionPasswordNeededError, PhoneCodeInvalidError
from telethon.tl.types import User

from config import API_ID, API_HASH, SESSION_DIR


class TelegramClientWrapper:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.session_file = os.path.join(SESSION_DIR, f"{session_id}.session")
        self.client: Optional[TelegramClient] = None
        self.phone: Optional[str] = None
        self.phone_code_hash: Optional[str] = None
        self._message_callback: Optional[Callable] = None

    async def init_client(self) -> TelegramClient:
        """Initialize or restore Telegram client"""
        os.makedirs(SESSION_DIR, exist_ok=True)

        # Use file-based session for persistence
        self.client = TelegramClient(
            self.session_file,
            API_ID,
            API_HASH,
            system_version="4.16.30-vxCUSTOM"
        )
        await self.client.connect()
        return self.client

    async def is_authorized(self) -> bool:
        """Check if client is authorized"""
        if not self.client:
            await self.init_client()
        return await self.client.is_user_authorized()

    async def send_code(self, phone: str) -> str:
        """Send verification code to phone"""
        if not self.client:
            await self.init_client()

        self.phone = phone
        result = await self.client.send_code_request(phone)
        self.phone_code_hash = result.phone_code_hash
        return self.phone_code_hash

    async def sign_in_with_code(self, phone: str, code: str, phone_code_hash: str) -> dict:
        """Sign in with verification code"""
        if not self.client:
            await self.init_client()

        try:
            user = await self.client.sign_in(phone, code, phone_code_hash=phone_code_hash)
            return {
                "success": True,
                "user_id": user.id,
                "username": user.username,
                "first_name": user.first_name
            }
        except SessionPasswordNeededError:
            return {
                "success": False,
                "needs_password": True,
                "message": "Two-factor authentication required"
            }
        except PhoneCodeInvalidError:
            return {
                "success": False,
                "needs_password": False,
                "message": "Invalid verification code"
            }

    async def sign_in_with_password(self, password: str) -> dict:
        """Sign in with 2FA password"""
        if not self.client:
            await self.init_client()

        try:
            user = await self.client.sign_in(password=password)
            return {
                "success": True,
                "user_id": user.id,
                "username": user.username,
                "first_name": user.first_name
            }
        except Exception as e:
            return {
                "success": False,
                "message": str(e)
            }

    async def get_me(self) -> Optional[User]:
        """Get current user info"""
        if not self.client or not await self.is_authorized():
            return None
        return await self.client.get_me()

    async def disconnect(self):
        """Disconnect client"""
        if self.client:
            await self.client.disconnect()