"""Concurrent download/convert/verify pipeline using asyncio.Queue.

Flow: download_queue → verify_queue ─┬─→ convert_queue → convert_verify_queue
                                      └─→ (done, if no conversion)

Each queue has dedicated workers with independent concurrency control.
"""

from __future__ import annotations

import asyncio
import errno
import json
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Awaitable

import aiohttp
import structlog

from xc_dl.core.manifest import Manifest, ManifestEntry
from xc_dl.core.progress import DownloadState
from xc_dl.core.verifier import compute_sha256, verify_content_length
from xc_dl.core.converter import convert_one
from xc_dl.utils import build_filename, build_relative_dir

logger = structlog.get_logger()

# Sentinel to signal workers to stop
_SENTINEL = None

# Queue item type aliases
DownloadItem = ManifestEntry
VerifyItem = tuple[ManifestEntry, Path, int]          # (entry, dest_path, file_size)
ConvertItem = ManifestEntry
ConvertVerifyItem = tuple[ManifestEntry, Path, str]   # (entry, converted_path, resampled_rel_path)


@dataclass
class PipelineStats:
    """Aggregate counters for a pipeline run."""
    download_ok: int = 0
    download_failed: int = 0
    download_skipped: int = 0
    convert_ok: int = 0
    convert_failed: int = 0
    convert_skipped: int = 0
    verify_ok: int = 0
    verify_failed: int = 0
    bytes_downloaded: int = 0


# Callback type aliases
ProgressCallback = Callable[..., Awaitable[None]] | None


class DownloadPipeline:
    """Queue-based pipeline: download → verify → convert → convert-verify.

    Items flow through stages as they complete, allowing concurrent
    operation across all four phases.
    """

    def __init__(
        self,
        session: aiohttp.ClientSession,
        data_dir: Path,
        cancel_event: asyncio.Event,
        concurrency: int = 8,
        max_retries: int = 3,
        skip_existing: bool = True,
        retry_failed: bool = True,
        convert_format: str | None = None,
        convert_sample_rate: int = 16000,
        convert_channels: int = 1,
        convert_bit_depth: int = 16,
        max_storage_bytes: int | None = None,
        min_free_space_bytes: int = 1_073_741_824,
        node_suffix: str = "",
        # Per-file progress callbacks
        on_download_start: ProgressCallback = None,
        on_download_progress: ProgressCallback = None,
        on_download_done: ProgressCallback = None,
        on_verify_start: ProgressCallback = None,
        on_verify_done: ProgressCallback = None,
        on_convert_start: ProgressCallback = None,
        on_convert_done: ProgressCallback = None,
        on_convert_verify_start: ProgressCallback = None,
        on_convert_verify_done: ProgressCallback = None,
        # Overall counters
        on_overall_download: ProgressCallback = None,
        on_overall_verify: ProgressCallback = None,
        on_overall_convert: ProgressCallback = None,
        on_overall_convert_verify: ProgressCallback = None,
    ):
        self.session = session
        self.data_dir = data_dir
        self.cancel_event = cancel_event
        self.concurrency = concurrency
        self.max_retries = max_retries
        self.skip_existing = skip_existing
        self.retry_failed = retry_failed
        self.convert_format = convert_format
        self.convert_sample_rate = convert_sample_rate
        self.convert_channels = convert_channels
        self.convert_bit_depth = convert_bit_depth
        self.max_storage_bytes = max_storage_bytes
        self.min_free_space_bytes = min_free_space_bytes
        self._space_check_counter = 0
        self._space_check_interval = max(concurrency, 8)

        # Queues
        self.download_queue: asyncio.Queue[ManifestEntry | None] = asyncio.Queue(
            maxsize=concurrency * 4
        )
        self.verify_queue: asyncio.Queue[VerifyItem | None] = asyncio.Queue()
        self.convert_queue: asyncio.Queue[ManifestEntry | None] = asyncio.Queue()
        self.convert_verify_queue: asyncio.Queue[ConvertVerifyItem | None] = asyncio.Queue()

        # Concurrency controls
        self.download_sem = asyncio.Semaphore(concurrency)
        self.verify_sem = asyncio.Semaphore(max(concurrency, 4))
        self.convert_sem = asyncio.Semaphore(max(concurrency // 2, 2))
        self.convert_verify_sem = asyncio.Semaphore(max(concurrency, 4))

        # Worker counts
        self.verify_workers_count = max(concurrency, 4)
        self.convert_workers_count = max(concurrency // 2, 2) if convert_format else 0
        self.convert_verify_workers_count = max(concurrency, 4) if convert_format else 0

        progress_dir = data_dir / ".progress"
        progress_dir.mkdir(parents=True, exist_ok=True)
        self.download_state = DownloadState(progress_dir, suffix=node_suffix)
        self.manifest = Manifest(data_dir)

        self.stats = PipelineStats()
        self._stats_lock = asyncio.Lock()

        # Callbacks
        self.on_download_start = on_download_start
        self.on_download_progress = on_download_progress
        self.on_download_done = on_download_done
        self.on_verify_start = on_verify_start
        self.on_verify_done = on_verify_done
        self.on_convert_start = on_convert_start
        self.on_convert_done = on_convert_done
        self.on_convert_verify_start = on_convert_verify_start
        self.on_convert_verify_done = on_convert_verify_done
        self.on_overall_download = on_overall_download
        self.on_overall_verify = on_overall_verify
        self.on_overall_convert = on_overall_convert
        self.on_overall_convert_verify = on_overall_convert_verify

    async def run(self, entries: list[ManifestEntry]) -> PipelineStats:
        """Run the full pipeline on the given entries."""
        # Filter entries that need work
        to_download: list[ManifestEntry] = []
        for entry in entries:
            status = self.download_state.get_status(entry.xc_id)
            if status == "verified" and self.skip_existing:
                if (self.data_dir / entry.original_path).exists():
                    async with self._stats_lock:
                        self.stats.download_skipped += 1
                    if self.on_overall_download:
                        await self.on_overall_download()
                    if self.on_overall_verify:
                        await self.on_overall_verify()
                    continue
            if status == "failed" and not self.retry_failed:
                async with self._stats_lock:
                    self.stats.download_skipped += 1
                if self.on_overall_download:
                    await self.on_overall_download()
                if self.on_overall_verify:
                    await self.on_overall_verify()
                continue
            to_download.append(entry)

        logger.info(
            "pipeline_plan",
            total=len(entries),
            to_download=len(to_download),
            skipped=self.stats.download_skipped,
        )

        if not to_download:
            return self.stats

        # Pre-flight disk space check
        try:
            usage = shutil.disk_usage(self.data_dir)
            if usage.free < self.min_free_space_bytes:
                logger.error(
                    "insufficient_disk_space",
                    free_bytes=usage.free,
                    min_free_bytes=self.min_free_space_bytes,
                )
                return self.stats
        except OSError:
            pass  # If we can't check, proceed and let downloads fail naturally

        # Enqueue all entries
        for entry in to_download:
            await self.download_queue.put(entry)

        # Add sentinels for download workers
        for _ in range(self.concurrency):
            await self.download_queue.put(_SENTINEL)

        # Start all workers
        download_tasks = [
            asyncio.create_task(self._download_worker(i))
            for i in range(self.concurrency)
        ]
        verify_tasks = [
            asyncio.create_task(self._verify_worker(i))
            for i in range(self.verify_workers_count)
        ]
        convert_tasks = [
            asyncio.create_task(self._convert_worker(i))
            for i in range(self.convert_workers_count)
        ]
        cv_tasks = [
            asyncio.create_task(self._convert_verify_worker(i))
            for i in range(self.convert_verify_workers_count)
        ]

        # Sentinel cascade: download → verify → convert → convert_verify
        await asyncio.gather(*download_tasks)

        for _ in range(self.verify_workers_count):
            await self.verify_queue.put(_SENTINEL)
        await asyncio.gather(*verify_tasks)

        for _ in range(self.convert_workers_count):
            await self.convert_queue.put(_SENTINEL)
        if convert_tasks:
            await asyncio.gather(*convert_tasks)

        for _ in range(self.convert_verify_workers_count):
            await self.convert_verify_queue.put(_SENTINEL)
        if cv_tasks:
            await asyncio.gather(*cv_tasks)

        # Update manifest for all downloaded entries
        for entry in to_download:
            status = self.download_state.get_status(entry.xc_id)
            if status == "verified":
                state = self.download_state._state.get(entry.xc_id, {})
                self.manifest.update_download_status(
                    entry.xc_id,
                    sha256=state.get("sha256", ""),
                    file_size=state.get("file_size_bytes", 0),
                )

        return self.stats

    async def _download_worker(self, worker_id: int) -> None:
        """Download worker: pulls from download_queue, pushes to verify_queue."""
        while not self.cancel_event.is_set():
            try:
                entry = await asyncio.wait_for(self.download_queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                if self.cancel_event.is_set():
                    break
                continue

            if entry is _SENTINEL:
                self.download_queue.task_done()
                break

            # Check storage cap
            if self.max_storage_bytes and self.stats.bytes_downloaded >= self.max_storage_bytes:
                logger.warning(
                    "storage_cap_reached",
                    downloaded=self.stats.bytes_downloaded,
                    cap=self.max_storage_bytes,
                )
                self.cancel_event.set()
                self.download_queue.task_done()
                break

            ok = await self._download_one(entry, worker_id)

            if ok:
                # Periodic disk space check after successful downloads
                async with self._stats_lock:
                    self._space_check_counter += 1
                    check_now = self._space_check_counter % self._space_check_interval == 0
                if check_now:
                    try:
                        usage = shutil.disk_usage(self.data_dir)
                        if usage.free < self.min_free_space_bytes:
                            logger.warning(
                                "low_disk_space",
                                free_bytes=usage.free,
                                min_free_bytes=self.min_free_space_bytes,
                            )
                            self.cancel_event.set()
                    except OSError:
                        pass
            else:
                # Failed downloads still advance overall download counter
                if self.on_overall_download:
                    await self.on_overall_download()

            self.download_queue.task_done()

    async def _download_one(self, entry: ManifestEntry, worker_id: int) -> bool:
        """Download a single file with retries. Content-length check inline, SHA256 deferred to verify_queue."""
        dest = self.data_dir / entry.original_path
        dest.parent.mkdir(parents=True, exist_ok=True)

        for attempt in range(1, self.max_retries + 1):
            if self.cancel_event.is_set():
                return False

            retry_delay = 0.0
            success = False

            async with self.download_sem:
                task_id = None
                self.download_state.mark_downloading(entry.xc_id)
                try:
                    timeout = aiohttp.ClientTimeout(total=300, sock_read=60)
                    async with self.session.get(entry.download_url, timeout=timeout) as resp:
                        if resp.status != 200:
                            logger.warning(
                                "download_http_error",
                                xc_id=entry.xc_id,
                                status=resp.status,
                                attempt=attempt,
                            )
                            if resp.status in (429, 503) and attempt < self.max_retries:
                                retry_delay = min(2**attempt, 60)
                            elif attempt >= self.max_retries:
                                self.download_state.mark_failed(entry.xc_id, f"HTTP {resp.status}", attempt)
                                async with self._stats_lock:
                                    self.stats.download_failed += 1
                                return False
                            # Non-retryable status on non-final attempt: fall through to retry
                        else:
                            expected_size = resp.content_length
                            received = 0

                            # Notify per-file progress
                            if self.on_download_start:
                                task_id = await self.on_download_start(
                                    entry.xc_id, entry.original_filename, expected_size
                                )

                            tmp = dest.with_suffix(dest.suffix + ".tmp")
                            cancelled = False
                            with open(tmp, "wb") as f:
                                async for chunk in resp.content.iter_chunked(1 << 16):
                                    f.write(chunk)
                                    received += len(chunk)
                                    if self.on_download_progress and task_id is not None:
                                        await self.on_download_progress(task_id, received)
                                    if self.cancel_event.is_set():
                                        cancelled = True
                                        break

                            if cancelled:
                                tmp.unlink(missing_ok=True)
                                return False

                            # Content-length check (instant) — retry on mismatch
                            if not verify_content_length(tmp, expected_size):
                                logger.warning("download_verify_failed", xc_id=entry.xc_id, attempt=attempt)
                                tmp.unlink(missing_ok=True)
                                if attempt >= self.max_retries:
                                    self.download_state.mark_failed(entry.xc_id, "content_length_mismatch", attempt)
                                    async with self._stats_lock:
                                        self.stats.download_failed += 1
                                    return False
                                # Fall through to retry
                            else:
                                # Content-length OK — rename to final destination
                                tmp.rename(dest)
                                file_size = dest.stat().st_size

                                async with self._stats_lock:
                                    self.stats.download_ok += 1
                                    self.stats.bytes_downloaded += file_size

                                logger.info(
                                    "download_complete",
                                    xc_id=entry.xc_id,
                                    stage="download",
                                    size_bytes=file_size,
                                )

                                # Advance overall download counter
                                if self.on_overall_download:
                                    await self.on_overall_download()

                                # Push to verify queue for SHA256 computation
                                await self.verify_queue.put((entry, dest, file_size))
                                success = True

                except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as exc:
                    # ENOSPC: clean up and stop immediately (no retries)
                    if isinstance(exc, OSError) and exc.errno == errno.ENOSPC:
                        logger.error("disk_full", xc_id=entry.xc_id, error=str(exc))
                        tmp = dest.with_suffix(dest.suffix + ".tmp")
                        tmp.unlink(missing_ok=True)
                        self.cancel_event.set()
                        async with self._stats_lock:
                            self.stats.download_failed += 1
                        return False
                    logger.warning("download_error", xc_id=entry.xc_id, error=str(exc), attempt=attempt)
                    if attempt >= self.max_retries:
                        self.download_state.mark_failed(entry.xc_id, str(exc), attempt)
                        async with self._stats_lock:
                            self.stats.download_failed += 1
                        return False
                    retry_delay = min(2**attempt, 60)
                finally:
                    if self.on_download_done and task_id is not None:
                        await self.on_download_done(task_id)

            if success:
                return True

            # Sleep OUTSIDE semaphore — other workers can proceed
            if retry_delay > 0:
                await asyncio.sleep(retry_delay)

        return False

    async def _verify_worker(self, worker_id: int) -> None:
        """Verify worker: computes SHA256 for downloaded files, pushes to convert_queue."""
        while not self.cancel_event.is_set():
            try:
                item = await asyncio.wait_for(self.verify_queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                if self.cancel_event.is_set():
                    break
                continue

            if item is _SENTINEL:
                self.verify_queue.task_done()
                break

            entry, dest, file_size = item

            async with self.verify_sem:
                task_id = None
                if self.on_verify_start:
                    task_id = await self.on_verify_start(entry.xc_id, entry.original_filename)

                sha256 = await compute_sha256(dest)
                self.download_state.mark_verified(entry.xc_id, sha256, file_size)

                async with self._stats_lock:
                    self.stats.verify_ok += 1

                logger.info(
                    "verify_complete",
                    xc_id=entry.xc_id,
                    stage="verify",
                    sha256=sha256,
                )

                if self.on_verify_done and task_id is not None:
                    await self.on_verify_done(task_id)

            if self.on_overall_verify:
                await self.on_overall_verify()

            # Push to convert queue if conversion requested
            if self.convert_format:
                await self.convert_queue.put(entry)

            self.verify_queue.task_done()

    async def _convert_worker(self, worker_id: int) -> None:
        """Convert worker: pulls from convert_queue, pushes to convert_verify_queue."""
        if not self.convert_format:
            return

        rate_label = f"{self.convert_sample_rate // 1000}khz"
        resampled_dir_name = f"resampled_{rate_label}"

        while not self.cancel_event.is_set():
            try:
                entry = await asyncio.wait_for(self.convert_queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                if self.cancel_event.is_set():
                    break
                continue

            if entry is _SENTINEL:
                self.convert_queue.task_done()
                break

            src = self.data_dir / entry.original_path
            if not src.exists():
                async with self._stats_lock:
                    self.stats.convert_failed += 1
                if self.on_overall_convert:
                    await self.on_overall_convert()
                self.convert_queue.task_done()
                continue

            rel_dir = build_relative_dir(entry.family, entry.gen, entry.sp)
            dst_name = build_filename(entry.xc_id, entry.gen, entry.sp, self.convert_format)
            dst = self.data_dir / resampled_dir_name / rel_dir / dst_name

            if dst.exists():
                async with self._stats_lock:
                    self.stats.convert_ok += 1
                if self.on_overall_convert:
                    await self.on_overall_convert()
                self.convert_queue.task_done()
                continue

            async with self.convert_sem:
                task_id = None
                if self.on_convert_start:
                    task_id = await self.on_convert_start(entry.xc_id, entry.original_filename)

                ok = await convert_one(
                    src, dst,
                    sample_rate=self.convert_sample_rate,
                    channels=self.convert_channels,
                    bit_depth=self.convert_bit_depth,
                )

                if self.on_convert_done and task_id is not None:
                    await self.on_convert_done(task_id)

            if ok:
                resampled_path = str(Path(resampled_dir_name) / rel_dir / dst_name)

                logger.info(
                    "convert_complete",
                    xc_id=entry.xc_id,
                    stage="convert",
                    dst=str(dst),
                )

                async with self._stats_lock:
                    self.stats.convert_ok += 1

                if self.on_overall_convert:
                    await self.on_overall_convert()

                # Push to convert-verify queue for SHA256
                await self.convert_verify_queue.put((entry, dst, resampled_path))
            else:
                async with self._stats_lock:
                    self.stats.convert_failed += 1
                if self.on_overall_convert:
                    await self.on_overall_convert()
                # Check if conversion failed due to low disk space
                try:
                    usage = shutil.disk_usage(self.data_dir)
                    if usage.free < self.min_free_space_bytes:
                        logger.warning(
                            "low_disk_space_convert",
                            free_bytes=usage.free,
                            min_free_bytes=self.min_free_space_bytes,
                        )
                        self.cancel_event.set()
                except OSError:
                    pass

            self.convert_queue.task_done()

    async def _convert_verify_worker(self, worker_id: int) -> None:
        """Convert-verify worker: computes SHA256 for converted files, updates manifest."""
        while not self.cancel_event.is_set():
            try:
                item = await asyncio.wait_for(self.convert_verify_queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                if self.cancel_event.is_set():
                    break
                continue

            if item is _SENTINEL:
                self.convert_verify_queue.task_done()
                break

            entry, dst, resampled_path = item

            async with self.convert_verify_sem:
                task_id = None
                if self.on_convert_verify_start:
                    task_id = await self.on_convert_verify_start(entry.xc_id, dst.name)

                sha256 = await compute_sha256(dst)
                self.manifest.update_entry(
                    entry.xc_id,
                    resampled_path=resampled_path,
                    sha256_resampled=sha256,
                    converted_from=entry.original_format,
                    sample_rate=self.convert_sample_rate,
                    channels=self.convert_channels,
                )

                # Update sidecar JSON with conversion info
                sidecar_path = self.data_dir / entry.metadata_path
                if sidecar_path.exists():
                    try:
                        sidecar = json.loads(sidecar_path.read_text())
                        xc_dl = sidecar.get("xc_dl", {})
                        xc_dl["resampled_path"] = resampled_path
                        xc_dl["sha256_resampled"] = sha256
                        xc_dl["converted_to"] = self.convert_format
                        sidecar["xc_dl"] = xc_dl
                        sidecar_path.write_text(json.dumps(sidecar, indent=2, ensure_ascii=False))
                    except (json.JSONDecodeError, OSError):
                        logger.warning("sidecar_update_failed", xc_id=entry.xc_id)

                logger.info(
                    "convert_verify_complete",
                    xc_id=entry.xc_id,
                    stage="convert_verify",
                    sha256=sha256,
                )

                if self.on_convert_verify_done and task_id is not None:
                    await self.on_convert_verify_done(task_id)

            if self.on_overall_convert_verify:
                await self.on_overall_convert_verify()

            self.convert_verify_queue.task_done()
