from __future__ import annotations

import asyncio
import os
import tempfile
from collections.abc import Callable
from pathlib import Path

import httpx

from .core import build_segments, detect_range_support, retry_async
from .models import DownloadConfig, DownloadProgress, DownloadResult, DownloadTask

ProgressCallback = Callable[[DownloadProgress], None]


class Downloader:
    _global_limit = 4
    _global_semaphore = asyncio.Semaphore(_global_limit)

    @classmethod
    def set_global_concurrency_limit(cls, limit: int) -> None:
        if limit < 1:
            raise ValueError("Global concurrency limit must be >= 1")
        cls._global_limit = limit
        cls._global_semaphore = asyncio.Semaphore(limit)

    def __init__(self, config: DownloadConfig | None = None) -> None:
        self.config = config or DownloadConfig()
        self._instance_lock = asyncio.Lock()
        self._is_downloading = False

    async def download(
        self,
        task: DownloadTask,
        progress_callback: ProgressCallback | None = None,
        headers: dict[str, str] | None = None,
    ) -> DownloadResult:
        destination = Path(task.output_path)
        destination.parent.mkdir(parents=True, exist_ok=True)

        async with self._instance_lock:
            if self._is_downloading:
                raise RuntimeError("This Downloader instance is already downloading a file")
            self._is_downloading = True

        try:
            async with type(self)._global_semaphore:
                timeout = httpx.Timeout(self.config.timeout_seconds)
                async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
                    return await self._download_with_strategy(client, task, destination, progress_callback, headers)
        finally:
            async with self._instance_lock:
                self._is_downloading = False

    async def _download_with_strategy(
        self,
        client: httpx.AsyncClient,
        task: DownloadTask,
        output_path: Path,
        progress_callback: ProgressCallback | None,
        headers: dict[str, str] | None,
    ) -> DownloadResult:
        range_support = await detect_range_support(
            client=client,
            url=task.url,
            retry_policy=self.config.retry_policy,
            headers=headers,
        )
        task.support_range = range_support.supported
        task.content_length = range_support.content_length

        use_multi = (
            range_support.supported
            and range_support.content_length is not None
            and range_support.content_length > 0
            and self.config.workers > 1
        )

        if use_multi and range_support.content_length is not None:
            segments = build_segments(range_support.content_length, self.config.workers)
            task.download_plan = {
                part_id: (start, end, end - start + 1)
                for part_id, (start, end) in enumerate(segments)
            }
            task.download_progress = {part_id: 0.0 for part_id in task.download_plan}
            task.total_downloaded = 0
            bytes_written = await self._download_multi_part(
                client,
                task,
                output_path,
                range_support.content_length,
                progress_callback,
                headers,
            )
        else:
            task.download_plan = None
            task.download_progress = None
            task.total_downloaded = 0
            bytes_written = await self._download_single(task, output_path, progress_callback, headers)

        return DownloadResult(
            url=task.url,
            output_path=output_path,
            bytes_written=bytes_written,
            used_multi_thread=use_multi,
            range_supported=range_support.supported,
        )

    async def _download_single(
        self,
        task: DownloadTask,
        output_path: Path,
        progress_callback: ProgressCallback | None,
        headers: dict[str, str] | None,
    ) -> int:
        request_headers = dict(headers or {})
        request_headers.pop("Range", None)

        async def action() -> httpx.Response:
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(self.config.timeout_seconds),
                follow_redirects=True,
            ) as client:
                return await client.get(task.url, headers=request_headers)

        response = await retry_async(action, self.config.retry_policy)
        downloaded = 0
        response.raise_for_status()
        content = response.content
        total = len(content)
        with output_path.open("wb") as file_handle:
            for offset in range(0, len(content), self.config.chunk_size):
                chunk = content[offset : offset + self.config.chunk_size]
                if not chunk:
                    continue
                file_handle.write(chunk)
                downloaded += len(chunk)
                task.total_downloaded = downloaded
                if progress_callback:
                    progress_callback(
                        DownloadProgress(
                            downloaded_bytes=downloaded,
                            total_bytes=total,
                            worker_count=1,
                            mode="single",
                        )
                    )
        await response.aclose()
        return downloaded

    async def _download_multi_part(
        self,
        client: httpx.AsyncClient,
        task: DownloadTask,
        output_path: Path,
        total_size: int,
        progress_callback: ProgressCallback | None,
        headers: dict[str, str] | None,
    ) -> int:
        segments = build_segments(total_size, self.config.workers)
        part_map = {part_id: segment for part_id, segment in enumerate(segments)}
        downloaded = 0
        downloaded_lock = asyncio.Lock()

        use_memory = self.config.mem_assemble

        part_data: dict[int, bytes] = {}
        part_files: dict[int, Path] = {}
        if use_memory:
            pass
        else:
            temp_dir = tempfile.gettempdir()
            part_files = {
                part_id: Path(temp_dir) / f".thread_downloader_{os.getpid()}_{part_id}.tmp"
                for part_id in part_map
            }

        async def worker(part_id: int, segment: tuple[int, int]) -> int:
            nonlocal downloaded
            start, end = segment
            request_headers = dict(headers or {})
            request_headers["Range"] = f"bytes={start}-{end}"

            async def action() -> httpx.Response:
                return await client.get(task.url, headers=request_headers)

            response = await retry_async(action, self.config.retry_policy)
            if response.status_code != 206:
                await response.aclose()
                raise RuntimeError("Range request expected HTTP 206 response")

            body = response.content
            local_written = 0

            if use_memory:
                part_data[part_id] = body
                local_written = len(body)
                async with downloaded_lock:
                    downloaded += local_written
                    task.total_downloaded = downloaded
                    if task.download_progress is not None:
                        task.download_progress[part_id] = float(local_written)
                    if progress_callback:
                        progress_callback(
                            DownloadProgress(
                                downloaded_bytes=downloaded,
                                total_bytes=total_size,
                                worker_count=len(segments),
                                mode="multi",
                            )
                        )
            else:
                with part_files[part_id].open("wb") as local_file:
                    for offset in range(0, len(body), self.config.chunk_size):
                        chunk = body[offset : offset + self.config.chunk_size]
                        if not chunk:
                            continue
                        local_file.write(chunk)
                        chunk_size = len(chunk)
                        local_written += chunk_size
                        async with downloaded_lock:
                            downloaded += chunk_size
                            task.total_downloaded = downloaded
                            if task.download_progress is not None:
                                task.download_progress[part_id] = float(local_written)
                            if progress_callback:
                                progress_callback(
                                    DownloadProgress(
                                        downloaded_bytes=downloaded,
                                        total_bytes=total_size,
                                        worker_count=len(segments),
                                        mode="multi",
                                    )
                                )

            await response.aclose()
            return local_written

        tasks = [asyncio.create_task(worker(part_id, segment)) for part_id, segment in part_map.items()]
        try:
            await asyncio.gather(*tasks)

            if use_memory:
                with output_path.open("wb") as destination:
                    for part_id in sorted(part_data.keys()):
                        destination.write(part_data[part_id])
            else:
                with output_path.open("wb") as destination:
                    for part_id in sorted(part_files.keys()):
                        with part_files[part_id].open("rb") as part_handle:
                            destination.write(part_handle.read())
                for path in part_files.values():
                    if path.exists():
                        path.unlink()
        except Exception:
            if not use_memory:
                for path in part_files.values():
                    if path.exists():
                        path.unlink()
            raise

        return downloaded
