"""Document downloader utilities for URL lists and DataFrames."""

from __future__ import annotations

import hashlib
import logging
import mimetypes
import os
import pathlib
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Tuple, Union

import pandas as pd
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from ..fhir.auth import Auth

logger = logging.getLogger(__name__)


class DocumentDownloadError(RuntimeError):
    """Raised when a document download fails."""


@dataclass
class DocumentResult:
    url: str
    download_id: str
    path: str
    status_code: Optional[int]
    content_type: Optional[str] = None
    sha256: Optional[str] = None
    error: Optional[str] = None


class DocumentDownloader:
    """Download documents from URLs into a deterministic folder structure.

    Parameters
    ----------
    auth:
        Optional Auth or requests.Session to reuse.
    base_url:
        Optional base URL to prefix relative paths.
    output_dir:
        Parent folder where documents are stored.
    timeout:
        Request timeout in seconds.
    skip_existing:
        Skip downloads if the target file already exists.
    log_downloads:
        If True, logs each download.
    """

    def __init__(
        self,
        auth: Optional[Union[Auth, requests.Session]] = None,
        base_url: Optional[str] = None,
        output_dir: Union[str, pathlib.Path] = "documents_out",
        timeout: int = 30,
        skip_existing: bool = True,
        log_downloads: bool = False,
        max_workers: int = 4,
        retries: Optional[Retry] = None,
        save_mode: str = "auto",
        force_extension: Optional[str] = None,
    ) -> None:
        """Initialize the document downloader."""
        self.base_url = base_url or os.environ.get("FHIR_BASE_URL")
        self.output_dir = pathlib.Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.timeout = timeout
        self.skip_existing = skip_existing
        self.log_downloads = log_downloads
        self.max_workers = max(1, max_workers)
        self.save_mode = save_mode
        self.force_extension = force_extension

        if isinstance(auth, Auth):
            self.session = auth.session
            self._close_session_on_exit = False
        elif isinstance(auth, requests.Session):
            self.session = auth
            self._close_session_on_exit = False
        else:
            if os.environ.get("FHIR_AUTH_URL") or os.environ.get("FHIR_REFRESH_URL"):
                self.session = Auth(
                    auth_method="env",
                    auth_url=os.environ.get("FHIR_AUTH_URL"),
                    refresh_url=os.environ.get("FHIR_REFRESH_URL"),
                ).session
                self._close_session_on_exit = False
            else:
                self.session = requests.Session()
                self._close_session_on_exit = True

        if retries is None:
            retries = Retry(
                total=3,
                backoff_factor=0.5,
                status_forcelist=[502, 503, 504],
                allowed_methods=["GET"],
                raise_on_status=False,
            )
        adapter = HTTPAdapter(max_retries=retries)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

    def __enter__(self) -> "DocumentDownloader":
        """Return self for context manager usage."""
        return self

    def close(self) -> None:
        """Close the underlying session if owned."""
        if self._close_session_on_exit:
            self.session.close()

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        """Close the session on context manager exit."""
        self.close()

    @staticmethod
    def download_id(url: str) -> str:
        """Return a deterministic download ID for a URL."""
        return hashlib.sha256(url.encode()).hexdigest()

    def download_url(self, url: str) -> DocumentResult:
        """Download a single URL into its deterministic folder."""
        full_url = self._resolve_url(url)
        download_id = self.download_id(full_url)
        target_dir = self.output_dir / download_id
        target_dir.mkdir(parents=True, exist_ok=True)

        filename = _filename_from_url(full_url)
        target_path = target_dir / filename

        if self.skip_existing and target_path.exists():
            if self.log_downloads:
                logger.info("Skipping existing document %s", target_path)
            return DocumentResult(full_url, download_id, str(target_path), None, None)

        try:
            if self.log_downloads:
                logger.info("Downloading %s", full_url)
            with self.session.get(full_url, stream=True, timeout=self.timeout) as resp:
                status = resp.status_code
                resp.raise_for_status()
                content_type = resp.headers.get("Content-Type")
                ext = _extension_from_content_type(content_type)
                target_path = _resolve_target_path(
                    target_dir=target_dir,
                    filename=filename,
                    ext=ext,
                    save_mode=self.save_mode,
                    force_extension=self.force_extension,
                )
                hasher = hashlib.sha256()
                with open(target_path, "wb") as f:
                    for chunk in resp.iter_content(chunk_size=1024 * 1024):
                        if chunk:
                            f.write(chunk)
                            hasher.update(chunk)
            return DocumentResult(
                full_url,
                download_id,
                str(target_path),
                status,
                content_type=content_type,
                sha256=hasher.hexdigest(),
                error=None,
            )
        except Exception as exc:
            logger.exception("Failed to download %s", full_url)
            return DocumentResult(full_url, download_id, str(target_dir), None, None, None, str(exc))

    def download_urls(
        self,
        urls: Iterable[str],
        results_csv: Optional[Union[str, pathlib.Path]] = None,
    ) -> pd.DataFrame:
        """Download a list of URLs and return results as a DataFrame."""
        url_list = list(urls)
        df = pd.DataFrame({"document_url": url_list})
        return self.download_from_dataframe(
            df=df,
            url_col="document_url",
            results_csv=results_csv,
        )

    def download_from_dataframe(
        self,
        df: pd.DataFrame,
        url_col: str = "document_url",
        results_csv: Optional[Union[str, pathlib.Path]] = None,
    ) -> pd.DataFrame:
        """Download all URLs from a DataFrame column."""
        if url_col not in df.columns:
            raise ValueError(f"Column '{url_col}' not found in DataFrame.")
        urls = df[url_col].astype(str).tolist()
        results: List[DocumentResult] = []
        if self.log_downloads:
            logger.info("Downloading %s documents", len(urls))
        if self.max_workers > 1:
            with ThreadPoolExecutor(max_workers=self.max_workers) as ex:
                futures = [ex.submit(self.download_url, url) for url in urls]
                for fut in as_completed(futures):
                    results.append(fut.result())
        else:
            for url in urls:
                results.append(self.download_url(url))
        out = pd.DataFrame([r.__dict__ for r in results])
        if results_csv:
            pathlib.Path(results_csv).parent.mkdir(parents=True, exist_ok=True)
            out.to_csv(results_csv, index=False)
        return out

    def download_from_df(
        self,
        df: pd.DataFrame,
        url_col: str = "document_url",
        results_csv: Optional[Union[str, pathlib.Path]] = None,
    ) -> pd.DataFrame:
        """Short alias for download_from_dataframe."""
        return self.download_from_dataframe(df=df, url_col=url_col, results_csv=results_csv)

    def _resolve_url(self, url: str) -> str:
        """Resolve relative URLs against base_url."""
        if url.startswith("http://") or url.startswith("https://"):
            return url
        if not self.base_url:
            raise ValueError("Relative URL provided but base_url is not set.")
        return self.base_url.rstrip("/") + "/" + url.lstrip("/")


def _filename_from_url(url: str) -> str:
    """Derive a filename from a URL path."""
    name = pathlib.Path(url.split("?")[0]).name
    return name or "document"


def _extension_from_content_type(content_type: Optional[str]) -> Optional[str]:
    """Map content-type to a file extension, if known."""
    if not content_type:
        return None
    ext = mimetypes.guess_extension(content_type.split(";")[0].strip())
    return ext


def _resolve_target_path(
    target_dir: pathlib.Path,
    filename: str,
    ext: Optional[str],
    save_mode: str,
    force_extension: Optional[str],
) -> pathlib.Path:
    """Determine the final target path based on save mode and extension."""
    # If user forces an extension, always use that.
    if force_extension:
        ext = force_extension if force_extension.startswith(".") else f".{force_extension}"
        return target_dir / (filename + ext)

    # Auto mode: prefer content-type extension if present.
    if save_mode == "auto":
        if ext and not filename.endswith(ext):
            return target_dir / (filename + ext)
        return target_dir / filename

    # Force to a given type.
    if save_mode in {"txt", "pdf"}:
        forced = f".{save_mode}"
        return target_dir / (filename + forced)

    # Default fallback
    return target_dir / filename
