from __future__ import annotations
import hashlib
import os
import tempfile
from dataclasses import dataclass
from datetime import datetime
from email.message import Message
from pathlib import PurePosixPath
from urllib.parse import unquote

import httpx

from govcrawler.settings import get_settings
from govcrawler.storage.filenames import dedupe_filename, safe_filename, with_best_extension
from govcrawler.storage.paths import build_reldir, to_os_path
from govcrawler.utils.url_norm import normalize_url

MAX_ATTACHMENT_BYTES = 200 * 1024 * 1024   # 200 MB
CHUNK = 64 * 1024
SAFE_EXT = {"pdf", "doc", "docx", "xls", "xlsx", "zip", "rar", "txt"}


@dataclass
class DownloadedAttachment:
    file_name: str
    file_ext: str
    size_bytes: int
    file_hash: str
    file_path: PurePosixPath
    source_url: str | None = None


def parse_disposition_filename(header: str | None) -> str | None:
    if not header:
        return None
    m = Message()
    m["content-disposition"] = header
    name = m.get_param("filename*", header="content-disposition") or m.get_param(
        "filename", header="content-disposition"
    )
    if not name:
        return None
    if isinstance(name, tuple):
        charset, _, raw = name
        try:
            # Python's email lib decodes %XX as latin-1 bytes; re-encode then decode
            # with the declared charset to recover the original (e.g. utf-8) text.
            if isinstance(raw, str):
                cs = (charset or "utf-8").lower() or "utf-8"
                try:
                    return raw.encode("latin-1").decode(cs)
                except Exception:
                    return unquote(raw)
            return raw
        except Exception:
            return raw
    if isinstance(name, str):
        decoded = unquote(name)
        # Some servers send an already URL-encoded value inside the regular
        # filename= parameter instead of RFC 5987 filename*=. Browsers decode
        # this for downloads; match that behavior for stored attachment names.
        return decoded
    return name


def _ext_of(name: str) -> str:
    _, dot, ext = name.rpartition(".")
    return ext.lower() if dot else ""


def _dedupe_filename(abs_dir, name: str, digest: str) -> str:
    """Return a safe filename, appending a short digest only on collision."""
    return dedupe_filename(abs_dir, name, digest)


def save_attachment_bytes(
    data: bytes,
    *,
    headers: dict | None = None,
    fallback_name: str = "attachment.bin",
    site: str,
    column: str,
    when: datetime,
    article_key: str,
    source_url: str | None = None,
) -> DownloadedAttachment:
    """Persist attachment bytes already fetched by a specialized client."""
    headers = headers or {}
    if len(data) > MAX_ATTACHMENT_BYTES:
        raise ValueError(f"attachment exceeds {MAX_ATTACHMENT_BYTES} bytes")
    settings = get_settings()
    reldir = build_reldir(site, column, when, "attachments")
    abs_dir = to_os_path(settings.data_dir, reldir)
    abs_dir.mkdir(parents=True, exist_ok=True)

    hasher = hashlib.sha256()
    hasher.update(data)
    digest = hasher.hexdigest()
    raw_name = (
        parse_disposition_filename(headers.get("content-disposition"))
        or fallback_name
        or "attachment.bin"
    )
    final_name = _dedupe_filename(abs_dir, raw_name, digest)
    final_abs = abs_dir / final_name
    try:
        final_abs.resolve().relative_to(abs_dir.resolve())
    except Exception:
        raise ValueError(f"path traversal blocked for filename={raw_name!r}")

    fd, tmp = tempfile.mkstemp(prefix=".dl_", dir=str(abs_dir))
    try:
        with os.fdopen(fd, "wb") as f:
            f.write(data)
        os.replace(tmp, final_abs)
    except Exception:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass
        raise

    return DownloadedAttachment(
        file_name=final_name,
        file_ext=_ext_of(final_name),
        size_bytes=len(data),
        file_hash=digest,
        file_path=reldir / final_name,
        source_url=source_url,
    )


def download_attachment(
    url: str,
    *,
    site: str,
    column: str,
    when: datetime,
    article_key: str,
    preferred_name: str | None = None,
    timeout_s: float | None = None,
) -> DownloadedAttachment:
    """Stream-download; enforce size cap + path-traversal guard."""
    norm = normalize_url(url)
    settings = get_settings()
    timeout = float(timeout_s if timeout_s is not None else settings.attachment_timeout_s)
    reldir = build_reldir(site, column, when, "attachments")
    abs_dir = to_os_path(settings.data_dir, reldir)
    abs_dir.mkdir(parents=True, exist_ok=True)

    hasher = hashlib.sha256()
    size = 0
    headers = {"User-Agent": settings.user_agent}
    # Inject cookies from the same pool that fetcher.http_client uses, so
    # ctct-shielded hosts (gdqy.gov.cn etc.) accept attachment downloads
    # after playwright primed them. Without this, /attachment/*.pdf URLs
    # under such hosts would return 412 challenge HTML and the bare
    # httpx.stream below would happily save it as a 'pdf' — but the file
    # would actually be the JS challenge page. Pulling cookies from the
    # pool keeps the request authenticated.
    try:
        from urllib.parse import urlparse as _up
        from govcrawler.cookies import get_default_store
        pool_cookies = get_default_store().get(_up(norm).netloc.lower()) or {}
    except Exception:
        pool_cookies = {}
    fd, tmp = tempfile.mkstemp(prefix=".dl_", dir=str(abs_dir))
    raw_name_fallback = norm.rsplit("/", 1)[-1].split("?", 1)[0] or "attachment.bin"
    raw_name: str = raw_name_fallback
    try:
        with httpx.stream(
            "GET", norm, headers=headers, cookies=pool_cookies,
            follow_redirects=True, timeout=timeout,
        ) as r:
            r.raise_for_status()
            disp_name = parse_disposition_filename(r.headers.get("content-disposition"))
            raw_name = (
                with_best_extension(
                    preferred_name,
                    disp_name,
                    raw_name_fallback,
                    r.headers.get("content-type"),
                )
                if preferred_name
                else (disp_name or raw_name_fallback)
            )
            with os.fdopen(fd, "wb") as f:
                for chunk in r.iter_bytes(chunk_size=CHUNK):
                    if not chunk:
                        continue
                    size += len(chunk)
                    if size > MAX_ATTACHMENT_BYTES:
                        raise ValueError(
                            f"attachment exceeds {MAX_ATTACHMENT_BYTES} bytes: {norm}"
                        )
                    hasher.update(chunk)
                    f.write(chunk)
    except Exception:
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass
        raise

    ext = _ext_of(raw_name)
    final_name = _dedupe_filename(abs_dir, raw_name, hasher.hexdigest())
    final_abs = abs_dir / final_name
    try:
        final_abs.resolve().relative_to(abs_dir.resolve())
    except Exception:
        os.unlink(tmp)
        raise ValueError(f"path traversal blocked for filename={raw_name!r}")
    os.replace(tmp, final_abs)

    return DownloadedAttachment(
        file_name=final_name,
        file_ext=ext,
        size_bytes=size,
        file_hash=hasher.hexdigest(),
        file_path=reldir / final_name,
        source_url=norm,
    )
