import asyncio, io, json, time
import aiohttp
from PIL import Image
from ..utils.logging import debug_log, log
from ..utils.network import get_client_session
from ..utils.constants import TILE_SEND_TIMEOUT
from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, _send_heartbeat_to_master
from ..utils.image import tensor_to_pil


class WorkerCommsMixin:
    async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id):
        """Proxy heartbeat helper used by worker processing mixins."""
        await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)

    async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, 
                                       padding, worker_id, is_final_flush=False):
        """Send all processed tiles to master, chunked if large."""
        if not processed_tiles:
            if is_final_flush:
                await self._send_tiles_completion_signal(multi_job_id, master_url, worker_id)
            return  # Early exit if empty

        total_tiles = len(processed_tiles)
        debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)")

        # Prepare encoded images and sizes to enable size-aware chunking
        encoded = []
        for idx, tile_data in enumerate(processed_tiles):
            img = tensor_to_pil(tile_data['tile'], 0)
            bio = io.BytesIO()
            # Keep compression low to balance speed and size; adjust if needed
            img.save(bio, format='PNG', compress_level=0)
            raw = bio.getvalue()
            encoded.append({
                'bytes': raw,
                'meta': {
                    'tile_idx': tile_data['tile_idx'],
                    'x': tile_data['x'],
                    'y': tile_data['y'],
                    'extracted_width': tile_data['extracted_width'],
                    'extracted_height': tile_data['extracted_height'],
                    **({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}),
                    **({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}),
                }
            })

        # Size-aware chunking
        max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024)  # 1MB headroom
        i = 0
        chunk_index = 0
        while i < total_tiles:
            data = aiohttp.FormData()
            data.add_field('multi_job_id', multi_job_id)
            data.add_field('worker_id', str(worker_id))
            data.add_field('padding', str(padding))

            metadata = []
            used = 0
            j = i
            while j < total_tiles:
                img_bytes = encoded[j]['bytes']
                meta = encoded[j]['meta']
                # Rough overhead for fields + JSON
                overhead = 1024
                if used + len(img_bytes) + overhead > max_bytes and j > i:
                    break
                # Accept this tile in this chunk
                metadata.append(meta)
                data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
                used += len(img_bytes) + overhead
                j += 1

            # Ensure at least one tile per chunk
            if j == i:
                # Single oversized tile, send anyway
                meta = encoded[j]['meta']
                metadata.append(meta)
                data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
                j += 1

            chunk_size = j - i
            is_chunk_last = (j >= total_tiles)
            data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
            data.add_field('batch_size', str(chunk_size))
            data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')

            # Retry logic with exponential backoff
            max_retries = 5
            retry_delay = 0.5
            for attempt in range(max_retries):
                try:
                    session = await get_client_session()
                    url = f"{master_url}/distributed/submit_tiles"
                    async with session.post(url, data=data) as response:
                        response.raise_for_status()
                        break
                except Exception as e:
                    if attempt < max_retries - 1:
                        await asyncio.sleep(retry_delay)
                        retry_delay = min(retry_delay * 2, 5.0)
                    else:
                        log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
                        raise

            debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
            chunk_index += 1
            i = j

    async def _send_tiles_completion_signal(self, multi_job_id, master_url, worker_id):
        """Send completion signal to master in static mode when no tiles are left."""
        data = aiohttp.FormData()
        data.add_field('multi_job_id', multi_job_id)
        data.add_field('worker_id', str(worker_id))
        data.add_field('is_last', 'true')
        data.add_field('batch_size', '0')

        session = await get_client_session()
        url = f"{master_url}/distributed/submit_tiles"
        async with session.post(url, data=data) as response:
            response.raise_for_status()
            debug_log(f"Worker {worker_id} sent static completion signal")

    async def _request_work_item_from_master(
        self,
        multi_job_id,
        master_url,
        worker_id,
        endpoint="/distributed/request_image",
    ):
        """Request one work item from master with retry/backoff and total timeout."""
        max_retries = 10
        retry_delay = 0.5
        start_time = time.monotonic()
        url = f"{master_url}{endpoint}"

        for attempt in range(max_retries):
            if time.monotonic() - start_time > 30:
                log(f"Total request timeout after 30s for worker {worker_id}")
                return None

            try:
                session = await get_client_session()
                async with session.post(url, json={
                    'worker_id': str(worker_id),
                    'multi_job_id': multi_job_id
                }) as response:
                    if response.status == 200:
                        return await response.json()
                    if response.status == 404:
                        text = await response.text()
                        debug_log(f"Job not found (404), will retry: {text}")
                        await asyncio.sleep(1.0)
                    else:
                        text = await response.text()
                        debug_log(
                            f"Request work item failed ({response.status}) for worker {worker_id}: {text}"
                        )

            except Exception as exc:
                if attempt < max_retries - 1:
                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {exc}")
                    await asyncio.sleep(retry_delay)
                    retry_delay = min(retry_delay * 2, 5.0)
                else:
                    log(f"Failed to request work item after {max_retries} attempts: {exc}")
                    raise

        return None

    async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
        """Request an image index to process from master in dynamic mode."""
        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
        if not data:
            return None, 0
        image_idx = data.get('image_idx')
        estimated_remaining = data.get('estimated_remaining', 0)
        return image_idx, estimated_remaining

    async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
        """Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
        if not data:
            return None, 0, False
        tile_idx = data.get('tile_idx')
        estimated_remaining = data.get('estimated_remaining', 0)
        batched_static = data.get('batched_static', False)
        return tile_idx, estimated_remaining, batched_static

    async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, 
                                        master_url, worker_id, is_last):
        """Send a processed full image back to master in dynamic mode."""
        # Serialize image to PNG
        byte_io = io.BytesIO()
        image_pil.save(byte_io, format='PNG', compress_level=0)
        byte_io.seek(0)
        
        # Prepare form data
        data = aiohttp.FormData()
        data.add_field('multi_job_id', multi_job_id)
        data.add_field('worker_id', str(worker_id))
        data.add_field('image_idx', str(image_idx))
        data.add_field('is_last', str(is_last))
        data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', 
                      content_type='image/png')
        
        # Retry logic
        max_retries = 5
        retry_delay = 0.5
        
        for attempt in range(max_retries):
            try:
                session = await get_client_session()
                url = f"{master_url}/distributed/submit_image"
                
                async with session.post(url, data=data) as response:
                    response.raise_for_status()
                    debug_log(f"Successfully sent image {image_idx} to master")
                    return
                    
            except Exception as e:
                if attempt < max_retries - 1:
                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
                    await asyncio.sleep(retry_delay)
                    retry_delay *= 2
                else:
                    log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
                    raise

    async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
        """Send completion signal to master in dynamic mode."""
        # Send a dummy request with is_last=True
        data = aiohttp.FormData()
        data.add_field('multi_job_id', multi_job_id)
        data.add_field('worker_id', str(worker_id))
        data.add_field('is_last', 'true')
        # No image data - just completion signal
        
        session = await get_client_session()
        url = f"{master_url}/distributed/submit_image"
        
        async with session.post(url, data=data) as response:
            response.raise_for_status()
            debug_log(f"Worker {worker_id} sent completion signal")

    async def _check_job_status(self, multi_job_id, master_url):
        """Check if job is ready on the master."""
        try:
            session = await get_client_session()
            url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
            async with session.get(url) as response:
                if response.status == 200:
                    data = await response.json()
                    return data.get('ready', False)
                return False
        except Exception as e:
            debug_log(f"Job status check failed: {e}")
            return False

    async def _async_yield(self):
        """Simple async yield to allow event loop processing."""
        await asyncio.sleep(0)