# 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 watchfiles import awatch

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


class ManagedClient:
    """
    Wrapper around Pyrogram Client
    Handles:
    - auto reconnect
    - per-client lock
    """

    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()


class MultiSessionRunner:
    """
    Production-grade client pool
    """

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

    # ==================================================
    # INITIAL LOAD
    # ==================================================

    async def load_all(self):
        async with self._registry_lock:
            await self._sync_with_disk()

        # start all once
        await self._start_all()

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

    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 to start {phone}: {e}")

    # ==================================================
    # SYNC WITH DISK (SAFE)
    # ==================================================

    async def _sync_with_disk(self):

        disk_phones = set()

        for file in os.listdir(SESSIONS_DIR):
            if file.endswith("_data.json"):
                phone = file.replace("_data.json", "")
                disk_phones.add(phone)

        current_phones = set(self.clients.keys())

        # ---- ADD NEW ----
        for phone in disk_phones - current_phones:
            client = self._create_client(phone)
            self.clients[phone] = ManagedClient(phone, client)
            print(f"[WATCHER] Added → {phone}")

        # ---- REMOVE ----
        for phone in current_phones - disk_phones:
            managed = self.clients.pop(phone)
            try:
                await managed.stop()
            except:
                pass
            print(f"[WATCHER] Removed → {phone}")

    # ==================================================
    # 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
        )

    # ==================================================
    # WATCHER (RACE SAFE)
    # ==================================================

    async def start_watcher(self):

        if self._watch_task:
            return

        async def watcher():
            async for _ in awatch(SESSIONS_DIR):
                async with self._registry_lock:
                    await self._sync_with_disk()

        self._watch_task = asyncio.create_task(watcher())
        print("🔥 Session watcher started")

    # ==================================================
    # 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 (NO SQLITE LOCK)
    # ==================================================

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

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