from __future__ import annotations

import hashlib
import json
import logging
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any, Callable

from niannian.storage.cloud_base import parse_remote_root_parts
from niannian.storage.upload_checkpoint import ChunkedUploadHooks, UploadCheckpointStore
from niannian.storage.upload_progress import UploadProgressTracker

logger = logging.getLogger(__name__)

DEFAULT_PCS_API_BASE = "https://pan.baidu.com/rest/2.0/xpan"
DEFAULT_PCS_SUPERFILE_BASE = "https://d.pcs.baidu.com/rest/2.0/pcs"
DEFAULT_UPLOAD_MAX_BYTES = 4 * 1024 * 1024
DEFAULT_CHUNK_SIZE_BYTES = 4 * 1024 * 1024
DEFAULT_MAX_UPLOAD_BYTES = 512 * 1024 * 1024
HttpOpener = Callable[..., Any]


def _split_chunks(data: bytes, chunk_size: int) -> list[bytes]:
    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive")
    return [data[i : i + chunk_size] for i in range(0, len(data), chunk_size)]


class BaiduPCSError(Exception):
    pass


def pcs_path(rel_path: str, remote_root: str) -> str:
    """rel_path: sync/manifest.json → /apps/念念/data/sync/manifest.json"""
    rel = rel_path.strip("/").replace("\\", "/")
    root = remote_root.strip("/")
    if not root:
        raise ValueError("remote_root must not be empty")
    if rel:
        return f"/{root}/{rel}"
    return f"/{root}"


class BaiduPCSClient:
    """百度 PCS 小文件 upload/download — Union OAuth access_token。"""

    def __init__(
        self,
        access_token: str,
        *,
        remote_root: str | None = None,
        pcs_api_base: str = DEFAULT_PCS_API_BASE,
        pcs_superfile_base: str = DEFAULT_PCS_SUPERFILE_BASE,
        upload_max_bytes: int = DEFAULT_UPLOAD_MAX_BYTES,
        chunk_size_bytes: int = DEFAULT_CHUNK_SIZE_BYTES,
        max_upload_bytes: int = DEFAULT_MAX_UPLOAD_BYTES,
        timeout: float = 30.0,
        http_opener: HttpOpener | None = None,
        state_path: Path | None = None,
        credential_manager: Any | None = None,
    ) -> None:
        self._access_token = access_token
        self._root_parts = parse_remote_root_parts(remote_root)
        self._remote_root = "/".join(self._root_parts)
        self._pcs_base = pcs_api_base.rstrip("/")
        self._superfile_base = pcs_superfile_base.rstrip("/")
        self._upload_max_bytes = int(upload_max_bytes)
        self._chunk_size_bytes = int(chunk_size_bytes)
        self._max_upload_bytes = int(max_upload_bytes)
        self._timeout = timeout
        self._open = http_opener or urllib.request.urlopen
        self._state_path = state_path
        self._credential_manager = credential_manager
        self._created_dirs: set[str] = set()
        self._load_state()
        self._prime_union_sandbox_dirs()

    def _union_sandbox_prefixes(self) -> set[str]:
        """Union OAuth apps get /apps/{product_name} pre-provisioned by Baidu."""
        if len(self._root_parts) >= 2 and self._root_parts[0] == "apps":
            return {"/apps", f"/apps/{self._root_parts[1]}"}
        return set()

    def _prime_union_sandbox_dirs(self) -> None:
        self._created_dirs.update(self._union_sandbox_prefixes())

    @property
    def access_token(self) -> str:
        return self._access_token

    def physical_path(self, rel_path: str) -> str:
        return pcs_path(rel_path, self._remote_root)

    def upload_object(
        self,
        rel_path: str,
        data: bytes,
        *,
        progress: UploadProgressTracker | None = None,
        checkpoint: UploadCheckpointStore | None = None,
        job_id: str | None = None,
    ) -> None:
        if len(data) > self._max_upload_bytes:
            logger.warning(
                "PCS upload rejected: %s exceeds max_upload_bytes (%d > %d)",
                rel_path,
                len(data),
                self._max_upload_bytes,
            )
            raise BaiduPCSError(
                f"file too large: {rel_path} ({len(data)} bytes > max {self._max_upload_bytes})"
            )
        path = self.physical_path(rel_path)
        self._ensure_remote_root()
        self._ensure_parent_dirs(path)
        if len(data) <= self._upload_max_bytes:
            self._upload_file(path, data)
            return
        hooks = ChunkedUploadHooks(
            provider="baidu",
            rel_path=rel_path,
            remote_root=self._remote_root,
            data=data,
            chunk_size=self._chunk_size_bytes,
            progress=progress,
            checkpoint=checkpoint,
            job_id=job_id,
        )
        if hooks.enabled:
            resume = hooks.begin()
            try:
                self._upload_file_chunked(path, data, hooks=hooks, resume=resume)
                hooks.done()
            except Exception as exc:
                hooks.failed(str(exc))
                raise
        else:
            self._upload_file_chunked(path, data)

    def _ensure_remote_root(self) -> None:
        """Ensure storage.cloud.remote_root exists on PCS before nested uploads."""
        root_path = f"/{self._remote_root}"
        if root_path in self._created_dirs:
            return
        parts = [p for p in self._remote_root.split("/") if p]
        start = 1
        if len(parts) >= 2 and parts[0] == "apps":
            start = min(len(parts), 3)
        for i in range(start, len(parts) + 1):
            dir_path = "/" + "/".join(parts[:i])
            if dir_path in self._created_dirs or self._should_skip_dir_create(dir_path):
                self._created_dirs.add(dir_path)
                continue
            if self._create_dir(dir_path):
                self._created_dirs.add(dir_path)
        self._save_state()

    def _should_skip_dir_create(self, dir_path: str) -> bool:
        return dir_path in self._union_sandbox_prefixes()

    def download_object(self, rel_path: str) -> bytes | None:
        path = self.physical_path(rel_path)
        if not self._file_exists(path):
            return None
        return self._download_file(path)

    def _load_state(self) -> None:
        if not self._state_path or not self._state_path.is_file():
            return
        try:
            raw = json.loads(self._state_path.read_text(encoding="utf-8"))
            if raw.get("remote_root") != self._remote_root:
                return
            self._created_dirs = set(raw.get("dirs") or [])
        except (OSError, json.JSONDecodeError):
            pass

    def _save_state(self) -> None:
        if not self._state_path:
            return
        self._state_path.parent.mkdir(parents=True, exist_ok=True)
        payload = {"remote_root": self._remote_root, "dirs": sorted(self._created_dirs)}
        self._state_path.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

    def _file_url(self, method: str, **params: str) -> str:
        query = {"method": method, "access_token": self._access_token}
        query.update(params)
        return f"{self._pcs_base}/file?{urllib.parse.urlencode(query, quote_via=urllib.parse.quote)}"

    def _post_form(
        self,
        method: str,
        form: dict[str, str],
        *,
        extra_query: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        query = {"method": method, "access_token": self._access_token}
        if extra_query:
            query.update(extra_query)
        url = f"{self._pcs_base}/file?{urllib.parse.urlencode(query, quote_via=urllib.parse.quote)}"
        body = urllib.parse.urlencode(form, quote_via=urllib.parse.quote).encode("utf-8")
        return self._request_json(
            url,
            method="POST",
            data=body,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )

    def _request(
        self,
        url: str,
        *,
        method: str = "GET",
        data: bytes | None = None,
        headers: dict[str, str] | None = None,
        auth_retry: bool = True,
    ) -> bytes:
        hdrs = dict(headers or {})
        if data is not None and "Content-Type" not in hdrs:
            hdrs.setdefault("Content-Type", "application/octet-stream")
        hdrs.setdefault("User-Agent", "niannian/0.2.8")
        req = urllib.request.Request(url, data=data, method=method, headers=hdrs)
        try:
            with self._open(req, timeout=self._timeout) as resp:
                status = getattr(resp, "status", 200)
                body = resp.read()
                if status >= 400:
                    raise urllib.error.HTTPError(
                        url, status, body.decode("utf-8", "replace"), resp.headers, None
                    )
                return body
        except urllib.error.HTTPError as exc:
            if auth_retry and exc.code == 401 and self._credential_manager is not None:
                refresh = self._credential_manager.refresh("baidu")
                if refresh.ok:
                    new_token = self._credential_manager.get_access_token("baidu")
                    if new_token:
                        self._access_token = new_token
                        parsed = urllib.parse.urlparse(url)
                        q = urllib.parse.parse_qs(parsed.query, keep_blank_values=True)
                        q["access_token"] = [new_token]
                        retry_url = urllib.parse.urlunparse(
                            parsed._replace(query=urllib.parse.urlencode(q, doseq=True))
                        )
                        return self._request(
                            retry_url,
                            method=method,
                            data=data,
                            headers=headers,
                            auth_retry=False,
                        )
            detail = exc.read().decode("utf-8", "replace") if exc.fp else str(exc)
            raise BaiduPCSError(f"PCS HTTP {exc.code}: {detail[:200]}") from exc

    def _request_json(
        self,
        url: str,
        *,
        method: str = "GET",
        data: bytes | None = None,
        headers: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        raw = self._request(url, method=method, data=data, headers=headers)
        try:
            payload = json.loads(raw.decode("utf-8", "replace"))
        except json.JSONDecodeError as exc:
            raise BaiduPCSError(f"invalid JSON from PCS: {raw[:200]!r}") from exc
        errno = payload.get("errno")
        if errno not in (0, None):
            raise BaiduPCSError(f"PCS errno={errno}: {payload.get('errmsg', payload)}")
        return payload

    def _ensure_parent_dirs(self, file_path: str) -> None:
        parts = [p for p in file_path.strip("/").split("/") if p]
        if len(parts) <= 1:
            return
        for i in range(1, len(parts)):
            dir_path = "/" + "/".join(parts[:i])
            if dir_path in self._created_dirs or self._should_skip_dir_create(dir_path):
                self._created_dirs.add(dir_path)
                continue
            if self._create_dir(dir_path):
                self._created_dirs.add(dir_path)
                self._save_state()

    def _create_dir(self, dir_path: str) -> bool:
        if not dir_path or dir_path == "/":
            return True
        try:
            self._post_form("create", {"path": dir_path, "isdir": "1", "rtype": "1"})
            return True
        except BaiduPCSError as exc:
            msg = str(exc)
            if "errno=-8" in msg or "errno=0" in msg or "already exists" in msg.lower():
                return True
            if "errno=2" in msg and dir_path.count("/") > 1:
                parent = "/" + "/".join(dir_path.strip("/").split("/")[:-1])
                if parent and parent not in self._created_dirs:
                    self._create_dir(parent)
                    self._created_dirs.add(parent)
                    self._post_form("create", {"path": dir_path, "isdir": "1", "rtype": "1"})
                    return True
            raise

    def _file_exists(self, path: str) -> bool:
        url = self._file_url("filemetas", dlink="0", path=path)
        try:
            payload = self._request_json(url, method="POST")
            metas = payload.get("list") or payload.get("info") or []
            return bool(metas)
        except BaiduPCSError:
            return False

    def _upload_file(self, path: str, data: bytes) -> None:
        """Union xpan 标准三步：precreate → superfile2 → create（非废弃 method=upload）。"""
        upload_id, block_md5s = self._precreate_and_upload_chunks(path, data, [data])
        self._finalize_upload(path, data, upload_id, block_md5s)

    def _upload_file_chunked(
        self,
        path: str,
        data: bytes,
        *,
        hooks: ChunkedUploadHooks | None = None,
        resume: Any | None = None,
    ) -> None:
        chunks = _split_chunks(data, self._chunk_size_bytes)
        upload_id, block_md5s = self._precreate_and_upload_chunks(
            path, data, chunks, hooks=hooks, resume=resume
        )
        self._finalize_upload(path, data, upload_id, block_md5s)

    def _precreate_and_upload_chunks(
        self,
        path: str,
        data: bytes,
        chunks: list[bytes],
        *,
        hooks: ChunkedUploadHooks | None = None,
        resume: Any | None = None,
    ) -> tuple[str, list[str]]:
        chunk_md5s = [hashlib.md5(chunk).hexdigest() for chunk in chunks]
        completed: set[int] = set()
        upload_id: str | None = None
        if resume is not None and resume.provider_state.get("uploadid"):
            upload_id = str(resume.provider_state["uploadid"])
            completed = set(resume.completed_chunks)
            stored_blocks = resume.provider_state.get("block_list")
            if isinstance(stored_blocks, list) and len(stored_blocks) == len(chunks):
                chunk_md5s = [str(x) for x in stored_blocks]

        if not upload_id:
            block_list = json.dumps(chunk_md5s, separators=(",", ":"))
            precreate = self._post_form(
                "precreate",
                {
                    "path": path,
                    "size": str(len(data)),
                    "isdir": "0",
                    "autoinit": "1",
                    "rtype": "3",
                    "block_list": block_list,
                },
            )
            upload_id = precreate.get("uploadid")
            if not upload_id:
                raise BaiduPCSError(f"precreate missing uploadid: {precreate}")

        file_name = path.rsplit("/", 1)[-1]
        result_md5s: list[str] = []
        for partseq, chunk in enumerate(chunks):
            if partseq in completed:
                result_md5s.append(chunk_md5s[partseq])
                continue
            body, content_type = self._multipart_body("file", file_name, chunk)
            query = {
                "method": "upload",
                "access_token": self._access_token,
                "type": "tmpfile",
                "path": path,
                "uploadid": upload_id,
                "partseq": str(partseq),
            }
            upload_url = (
                f"{self._superfile_base}/superfile2?"
                f"{urllib.parse.urlencode(query, quote_via=urllib.parse.quote)}"
            )
            upload_raw = self._request(
                upload_url,
                method="POST",
                data=body,
                headers={"Content-Type": content_type},
            )
            chunk_md5 = chunk_md5s[partseq]
            try:
                upload_payload = json.loads(upload_raw.decode("utf-8", "replace"))
                errno = upload_payload.get("errno")
                if errno is None and upload_payload.get("error_code") is not None:
                    errno = upload_payload.get("error_code")
                if errno not in (0, None):
                    detail = (
                        upload_payload.get("errmsg")
                        or upload_payload.get("error_msg")
                        or upload_payload
                    )
                    raise BaiduPCSError(f"superfile2 errno={errno}: {detail}")
                chunk_md5 = upload_payload.get("md5") or chunk_md5
            except json.JSONDecodeError:
                pass
            result_md5s.append(chunk_md5)
            if hooks is not None:
                hooks.after_chunk(
                    partseq,
                    provider_state={
                        "uploadid": upload_id,
                        "remote_path": path,
                        "block_list": chunk_md5s,
                        "completed_partseq": sorted(set(completed) | {partseq}),
                    },
                )
        return str(upload_id), result_md5s

    def _finalize_upload(
        self,
        path: str,
        data: bytes,
        upload_id: str,
        block_md5s: list[str],
    ) -> None:
        self._post_form(
            "create",
            {
                "path": path,
                "size": str(len(data)),
                "isdir": "0",
                "rtype": "1",
                "uploadid": upload_id,
                "block_list": json.dumps(block_md5s, separators=(",", ":")),
                "ondup": "overwrite",
            },
        )

    def _download_file(self, path: str) -> bytes:
        url = self._file_url("download", path=path)
        body = self._request(url, method="GET")
        if body.startswith(b"{") or body.startswith(b"["):
            payload = json.loads(body.decode("utf-8", "replace"))
            dlink = payload.get("dlink") or (payload.get("list") or [{}])[0].get("dlink")
            if dlink:
                return self._request(dlink, method="GET", headers={})
        return body

    def _multipart_body(self, field_name: str, filename: str, data: bytes) -> tuple[bytes, str]:
        boundary = "----niannianpcsupload"
        prefix = (
            f"--{boundary}\r\n"
            f"Content-Disposition: form-data; name=\"{field_name}\"; filename=\"{filename}\"\r\n"
            "Content-Type: application/octet-stream\r\n\r\n"
        ).encode()
        suffix = f"\r\n--{boundary}--\r\n".encode()
        return prefix + data + suffix, f"multipart/form-data; boundary={boundary}"
