import asyncio
import random
import uuid
import time
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Callable, Dict, Optional, Set


class JobStatus(Enum):
    QUEUED = auto()
    STARTED = auto()
    COMPLETED = auto()
    FAILED = auto()
    CANCELLED = auto()


@dataclass
class Job:
    id: str
    text: str
    voice: str
    status: JobStatus = JobStatus.QUEUED
    retries_attempted: int = 0
    cancelled: bool = False


async def mock_synthesize(text: str) -> bytes:
    """
    Simulates a TTS engine. 
    Sleeps 50-300ms, fails ~10% of the time.
    """
    await asyncio.sleep(random.uniform(0.05, 0.3))
    if random.random() < 0.1:
        raise Exception("Synthetic Synthesis Error")
    return b"\x00" * len(text)


class TTSPipeline:
    def __init__(self, max_concurrency: int = 4, max_queue_size: int = 100):
        self.max_concurrency = max_concurrency
        self.max_queue_size = max_queue_size
        
        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._jobs: Dict[str, Job] = {}
        self._workers: Set[asyncio.Task] = set()
        self._callback: Optional[Callable[[str, JobStatus], None]] = None
        
        self._semaphore = asyncio.Semaphore(max_concurrency)
        self._running_count = 0  # Tracks jobs in queue + jobs in flight
        self._lock = asyncio.Lock() # Protects shared state updates
        
        # Metrics for testing
        self._current_active_count = 0
        self._peak_concurrency = 0

    def set_callback(self, cb: Callable[[str, JobStatus], None]):
        self._callback = cb

    def _emit(self, job_id: str, status: JobStatus):
        if self._callback:
            self._callback(job_id, status)

    async def start(self):
        """Initialize the worker pool."""
        for _ in range(self.max_concurrency):
            worker = asyncio.create_task(self._worker())
            self._workers.add(worker)

    async def stop(self):
        """Gracefully shut down workers."""
        for worker in self._workers:
            worker.cancel()
        await asyncio.gather(*self._workers, return_exceptions=True)
        self._workers.clear()

    async def submit(self, text: str, voice: str) -> str:
        """Enqueue a job. Returns job_id immediately."""
        if self._queue.qsize() >= self.max_queue_size:
            raise BufferError("Pipeline queue is full (backpressure)")

        job_id = str(uuid.uuid4())
        job = Job(id=job_id, text=text, voice=voice)
        
        async with self._lock:
            self._jobs[job_id] = job
            self._running_count += 1
        
        await self._queue.put(job_id)
        self._emit(job_id, JobStatus.QUEUED)
        return job_id

    async def cancel(self, job_id: str):
        """Cancel a queued or in-flight job."""
        async with self._lock:
            if job_id in self._jobs:
                job = self._jobs[job_id]
                if job.status not in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
                    job.cancelled = True

    async def drain(self):
        """Wait until all queued and in-flight jobs are finished."""
        # 1. Wait for the queue to be empty (all items processed by workers)
        await self._queue.join()
        # 2. Wait for the running count to hit zero (all workers finished their logic)
        while True:
            async with self._lock:
                if self._running_count == 0:
                    break
            await asyncio.sleep(0.01)

    async def _worker(self):
        while True:
            job_id = await self._queue.get()
            try:
                async with self._lock:
                    job = self._jobs.get(job_id)
                
                if not job:
                    continue

                # Check if cancelled while sitting in queue
                if job.cancelled:
                    await self._finalize_job(job, JobStatus.CANCELLED)
                    continue

                # Process job with bounded concurrency and retries
                async with self._semaphore:
                    await self._process_with_retries(job)

            except asyncio.CancelledError:
                break
            finally:
                self._queue.task_done()

    async def _process_with_retries(self, job: Job):
        max_retries = 3
        backoff = 0.1

        for attempt in range(max_retries + 1):
            # Check cancellation before starting/retrying
            if job.cancelled:
                await self._finalize_job(job, JobStatus.CANCELLED)
                return

            try:
                # Track concurrency for testing purposes
                async with self._lock:
                    self._current_active_count += 1
                    self._peak_concurrency = max(self._peak_concurrency, self._current_active_count)
                
                self._emit(job.id, JobStatus.STARTED)
                
                # The actual work
                await mock_synthesize(job.text)

                # Check cancellation after the await (in case it was cancelled during synthesis)
                if job.cancelled:
                    await self._finalize_job(job, JobStatus.CANCELLED)
                else:
                    await self._finalize_job(job, JobStatus.COMPLETED)
                return

            except Exception as e:
                if job.cancelled:
                    await self._finalize_job(job, JobStatus.CANCELLED)
                    return
                
                if attempt < max_retries:
                    # Exponential backoff
                    await asyncio.sleep(backoff * (2 ** attempt))
                else:
                    # Exhausted retries
                    await self._finalize_job(job, JobStatus.FAILED)
            finally:
                async with self._lock:
                    self._current_active_count -= 1

    async def _finalize_job(self, job: Job, status: JobStatus):
        """Updates job state and decrements the global running counter."""
        async with self._lock:
            if job.status not in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
                job.status = status
                self._running_count -= 1
                self._emit(job.id, status)


# ==========================================
# TEST SUITE
# ==========================================

async def main():
    print("--- Starting TTS Pipeline Tests ---\n")
    pipeline = TTSPipeline(max_concurrency=4, max_queue_size=100)
    await pipeline.start()

    # Event tracking for assertions
    events = []
    def callback(job_id, status):
        events.append((job_id, status))

    pipeline.set_callback(callback)

    # --- Test A: Bounded Concurrency & Drain ---
    print("Test A: Submitting 50 jobs and checking concurrency...")
    job_ids = []
    for i in range(50):
        jid = await pipeline.submit(f"Text {i}", "en-US")
        job_ids.append(jid)

    await pipeline.drain()
    
    assert pipeline._peak_concurrency <= 4, f"Concurrency exceeded! Peak: {pipeline._peak_concurrency}"
    print(f"  [PASS] Max concurrency was {pipeline._peak_concurrency}/4")

    # --- Test B: Backpressure ---
    print("\nTest B: Testing backpressure (100 job cap)...")
    # Fill the queue to 100. Note: drain() cleared previous jobs, so current count is 0.
    # We submit 100 to fill it.
    for i in range(100):
        await pipeline.submit("Backpressure test", "en-US")
    
    try:
        await pipeline.submit("The breaking job", "en-US")
        print("  [FAIL] Pipeline accepted 101st job without error.")
    except BufferError:
        print("  [PASS] Pipeline correctly rejected 101st job with BufferError.")

    # Clear the queue for next test
    await pipeline.drain()

    # --- Test C: Cancellation & Recovery ---
    print("\nTest C: Mid-flight cancellation and recovery...")
    # Submit a batch
    batch_ids = []
    for i in range(10):
        batch_ids.append(await pipeline.submit("Cancel me", "en-US"))
    
    # Wait a tiny bit to ensure jobs are picked up by workers (in-flight)
    await asyncio.sleep(0.1)
    
    # Cancel the first job in the batch
    target_id = batch_ids[0]
    await pipeline.cancel(target_id)
    print(f"  Cancelled job: {target_id}")

    await pipeline.drain()
    
    # Verify the cancelled job is marked correctly in events
    cancelled_events = [e for e in events if e[0] == target_id and e[1] == JobStatus.CANCELLED]
    assert len(cancelled_events) > 0, "Job was not recorded as CANCELLED"
    print("  [PASS] Job cancellation detected.")

    # Verify pool is still usable
    print("  Verifying pool usability...")
    new_job = await pipeline.submit("Post-cancel job", "en-US")
    await pipeline.drain()
    
    # Check if the new job completed
    new_job_events = [e for e in events if e[0] == new_job and e[1] == JobStatus.COMPLETED]
    # Note: It might be FAILED if the 10% error hit, but it shouldn't be stuck.
    assert any(e[1] in (JobStatus.COMPLETED, JobStatus.FAILED) for e in events if e[0] == new_job), "Pool stuck!"
    print("  [PASS] Pool recovered and processed new jobs.")

    await pipeline.stop()
    print("\n--- All Tests Passed Successfully ---")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass