"""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 collections.abc import Awaitable
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable

import aiofiles
import aiohttp
import structlog

from xc_dl.core.converter import convert_one
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.utils import build_filename, build_relative_dir

logger = structlog.get_logger()


async def _update_sidecar(
    sidecar_path: Path,
    resampled_path: str,
    sha256: str,
    convert_format: str,
) -> None:
    """Read-modify-write a sidecar JSON file using async I/O."""
    async with aiofiles.open(sidecar_path) as f:
        sidecar = json.loads(await f.read())
    xc_dl = sidecar.get("xc_dl", {})
    xc_dl["resampled_path"] = resampled_path
    xc_dl["sha256_resampled"] = sha256
    xc_dl["converted_to"] = convert_format
    sidecar["xc_dl"] = xc_dl
    tmp = sidecar_path.with_suffix(".tmp")
    async with aiofiles.open(tmp, "w") as f:
        await f.write(json.dumps(sidecar, indent=2, ensure_ascii=False))
    tmp.rename(sidecar_path)


# 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 and (self.data_dir / entry.original_path).exists():
                logger.info("download_skipped", xc_id=entry.xc_id, reason="already_verified")
                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:
                logger.info("download_skipped", xc_id=entry.xc_id, reason="failed_no_retry")
                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

        async with self.manifest.writer() as mw:
            self._manifest_writer = mw
            return await self._run_pipeline(to_download)

    async def _run_pipeline(self, to_download: list[ManifestEntry]) -> PipelineStats:
        """Inner pipeline loop — runs inside the manifest writer context."""
        # 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

        # Start all workers BEFORE enqueuing so the bounded queue can drain
        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)
        ]

        # Feed entries into the bounded download queue from a separate task.
        # The queue is bounded (maxsize=concurrency*4), so `await put()` blocks
        # the feeder when full — providing backpressure without loading all
        # entries into memory at once.
        async def _feed_download_queue() -> None:
            for entry in to_download:
                await self.download_queue.put(entry)
            for _ in range(self.concurrency):
                await self.download_queue.put(_SENTINEL)

        feeder = asyncio.create_task(_feed_download_queue())

        # Sentinel cascade: download → verify → convert → convert_verify
        await asyncio.gather(*download_tasks)
        await feeder  # ensure feeder finished (should be done by now)

        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.submit_update(
                    self._manifest_writer,
                    entry.xc_id,
                    sha256_original=state.get("sha256", ""),
                    file_size_bytes=state.get("file_size_bytes", 0),
                    status="verified",
                    downloaded_at=datetime.now(timezone.utc).isoformat(),
                )

        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
                await 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:
                                await 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
                            async with aiofiles.open(tmp, "wb") as f:
                                async for chunk in resp.content.iter_chunked(1 << 16):
                                    await 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:
                                    await 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:
                        await 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)
                await 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():
                logger.error("convert_source_missing", xc_id=entry.xc_id, path=str(src))
                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():
                logger.info("convert_skipped_exists", xc_id=entry.xc_id, dst=str(dst))
                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:
                logger.error("convert_failed", xc_id=entry.xc_id, dst=str(dst))
                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.submit_update(
                    self._manifest_writer,
                    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:
                        await _update_sidecar(
                            sidecar_path,
                            resampled_path,
                            sha256,
                            self.convert_format,
                        )
                    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()
