# remottxrea/runner/multi_session_runner.py

import os
import asyncio
from typing import Dict

from pyrogram import Client
from pyrogram.errors import FloodWait, RPCError

from ..config.main_config import SESSIONS_DIR
from ..session.session_data_manager import session_data_manager
from ..core.delay_engine import delay_engine
from ..core.execution_tracker import ExecutionTracker


# ==========================================================
# MANAGED CLIENT
# ==========================================================

class ManagedClient:

    def __init__(self, phone: str, client: Client):
        self.phone = phone
        self.client = client
        self.lock = asyncio.Lock()
        self.started = False

    async def start(self):
        async with self.lock:
            if not self.started:
                await self.client.start()
                self.started = True

    async def stop(self):
        async with self.lock:
            if self.started:
                await self.client.stop()
                self.started = False

    async def ensure_connected(self):
        async with self.lock:
            if not self.client.is_connected:
                await self.client.connect()


# ==========================================================
# MULTI SESSION RUNNER (NO WATCHER VERSION)
# ==========================================================

class MultiSessionRunner:

    def __init__(self):
        self.clients: Dict[str, ManagedClient] = {}
        self._registry_lock = asyncio.Lock()

    # ==================================================
    # INITIAL LOAD (ONLY ON STARTUP)
    # ==================================================

    async def load_all(self):

        async with self._registry_lock:

            for file in os.listdir(SESSIONS_DIR):

                if not file.endswith("_data.json"):
                    continue

                phone = file.replace("_data.json", "")

                if phone in self.clients:
                    continue

                client = self._create_client(phone)
                self.clients[phone] = ManagedClient(phone, client)

        await self._start_all()

    # ==================================================
    # REGISTER NEW SESSION (🔥 مهم)
    # ==================================================

    async def register_session(self, phone: str):

        async with self._registry_lock:

            if phone in self.clients:
                return

            client = self._create_client(phone)
            managed = ManagedClient(phone, client)

            self.clients[phone] = managed

            await managed.start()

            print(f"[POOL] Registered → {phone}")

    # ==================================================
    # REMOVE SESSION
    # ==================================================

    async def remove_session(self, phone: str):

        async with self._registry_lock:

            managed = self.clients.pop(phone, None)

            if managed:
                await managed.stop()
                print(f"[POOL] Removed → {phone}")

    # ==================================================
    # START ALL
    # ==================================================

    async def _start_all(self):

        for phone, managed in self.clients.items():
            try:
                await managed.start()
                print(f"[POOL] Started → {phone}")
            except Exception as e:
                print(f"[POOL] Failed → {phone}: {e}")

    # ==================================================
    # CLIENT FACTORY
    # ==================================================

    def _create_client(self, phone):

        data = session_data_manager.load(phone)
        device = data["device"]

        return Client(
            name=os.path.join(SESSIONS_DIR, phone),
            api_id=device["api_id"],
            api_hash=device["api_hash"],
            device_model=device["device_model"],
            system_version=device["system_version"],
            app_version=device["app_version"],
            lang_code=device["lang_code"],
            workdir=SESSIONS_DIR,
            workers=1  # جلوگیری از race داخلی
        )

    # ==================================================
    # RUN ACTION SAFELY
    # ==================================================

    async def run_action(self, action_func, progress_message=None):

        tracker = ExecutionTracker()
        total = len(self.clients)
        tracker.start(total)

        for index, (phone, managed) in enumerate(
            self.clients.items(),
            start=1
        ):

            status = f"[{index}/{total}] {phone}"

            try:
                await managed.start()
                await managed.ensure_connected()

                async with managed.lock:

                    if progress_message:
                        await progress_message.edit_text(
                            f"{status} → Running"
                        )

                    await action_func(phone, managed.client)

                tracker.record(phone, "success")
                status_msg = f"{status} → Success"

            except FloodWait as e:
                tracker.record(phone, "flood", e)
                status_msg = f"{status} → Flood"

            except RPCError as e:
                tracker.record(phone, "rpc", e)
                status_msg = f"{status} → RPC"

            except Exception as e:
                tracker.record(phone, "failed", e)
                status_msg = f"{status} → Failed"

            print(status_msg)

            if progress_message:
                await progress_message.edit_text(status_msg)

            await delay_engine.wait()

        report = tracker.summary_text()

        if progress_message:
            await progress_message.edit_text(report)

        return tracker

    # ==================================================
    # SAFE LOGIN LOCK
    # ==================================================

    async def acquire_pool_lock(self):
        await self._registry_lock.acquire()

    def release_pool_lock(self):
        if self._registry_lock.locked():
            self._registry_lock.release()


# ==========================================================
# 🔥 GLOBAL RUNNER INSTANCE (همین لازم داشتی)
# ==========================================================

runner = MultiSessionRunner()