from __future__ import annotations

import hashlib
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from niannian.storage.manifest import utc_now_iso
from niannian.storage.upload_progress import UploadProgressTracker

logger = logging.getLogger(__name__)


def content_sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def make_checkpoint_id(provider: str, rel_path: str, data_sha: str, remote_root: str) -> str:
    raw = f"{provider}|{rel_path}|{data_sha}|{remote_root}"
    return "ckpt_" + hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]


@dataclass
class ResumeState:
    checkpoint_id: str | None = None
    job_id: str | None = None
    provider_state: dict[str, Any] = field(default_factory=dict)
    completed_chunks: list[int] = field(default_factory=list)
    start_chunk_index: int = 0


class UploadCheckpointStore:
    """Persistent multipart checkpoints under {home}/sync/upload_checkpoints/."""

    def __init__(
        self,
        home: Path,
        *,
        checkpoint_dir: str = "sync/upload_checkpoints",
        ttl_days: int = 7,
    ) -> None:
        self._dir = home / checkpoint_dir
        self._dir.mkdir(parents=True, exist_ok=True)
        self._ttl_days = int(ttl_days)

    @property
    def directory(self) -> Path:
        return self._dir

    def _path_for(self, checkpoint_id: str) -> Path:
        safe = checkpoint_id.replace("/", "_")
        return self._dir / f"{safe}.json"

    def _is_expired(self, payload: dict[str, Any]) -> bool:
        expires = payload.get("expires_at")
        if not expires:
            return False
        try:
            exp = datetime.fromisoformat(str(expires).replace("Z", "+00:00"))
            return exp <= datetime.now(timezone.utc)
        except ValueError:
            return False

    def purge_expired(self) -> int:
        removed = 0
        for path in self._dir.glob("*.json"):
            try:
                payload = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError):
                continue
            if self._is_expired(payload):
                try:
                    path.unlink()
                    removed += 1
                except OSError:
                    pass
        return removed

    def load(
        self,
        *,
        provider: str,
        rel_path: str,
        remote_root: str,
        data_sha: str,
    ) -> dict[str, Any] | None:
        self.purge_expired()
        ckpt_id = make_checkpoint_id(provider, rel_path, data_sha, remote_root)
        path = self._path_for(ckpt_id)
        if not path.is_file():
            return None
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return None
        if self._is_expired(payload):
            try:
                path.unlink()
            except OSError:
                pass
            return None
        if payload.get("content_sha256") != data_sha:
            return None
        return payload

    def save(self, payload: dict[str, Any]) -> None:
        ckpt_id = str(payload["checkpoint_id"])
        now = utc_now_iso()
        expires = (
            datetime.now(timezone.utc) + timedelta(days=self._ttl_days)
        ).replace(microsecond=0).isoformat().replace("+00:00", "Z")
        payload.setdefault("version", 1)
        payload["updated_at"] = now
        payload.setdefault("created_at", now)
        payload["expires_at"] = expires
        path = self._path_for(ckpt_id)
        path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    def delete(self, checkpoint_id: str) -> None:
        path = self._path_for(checkpoint_id)
        if path.is_file():
            try:
                path.unlink()
            except OSError:
                pass

    def list_checkpoints(self, *, provider: str | None = None) -> list[dict[str, Any]]:
        self.purge_expired()
        out: list[dict[str, Any]] = []
        for path in sorted(self._dir.glob("*.json")):
            try:
                payload = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError):
                continue
            if provider and payload.get("provider") != provider:
                continue
            out.append(payload)
        out.sort(key=lambda p: p.get("updated_at") or "", reverse=True)
        return out

    def clear(self, *, provider: str | None = None, all_providers: bool = False) -> int:
        removed = 0
        for path in self._dir.glob("*.json"):
            try:
                payload = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError):
                payload = {}
            if all_providers or provider is None or payload.get("provider") == provider:
                try:
                    path.unlink()
                    removed += 1
                except OSError:
                    pass
        return removed

    def count_pending(self) -> int:
        return len(self.list_checkpoints())


class ChunkedUploadHooks:
    """Orchestrate progress + checkpoint for multipart uploads (best-effort)."""

    def __init__(
        self,
        *,
        provider: str,
        rel_path: str,
        remote_root: str,
        data: bytes,
        chunk_size: int,
        progress: UploadProgressTracker | None = None,
        checkpoint: UploadCheckpointStore | None = None,
        job_id: str | None = None,
    ) -> None:
        self.provider = provider
        self.rel_path = rel_path
        self.remote_root = remote_root
        self.data = data
        self.chunk_size = chunk_size
        self.progress = progress
        self.checkpoint = checkpoint
        self.job_id = job_id
        self.data_sha = content_sha256(data)
        self.bytes_total = len(data)
        self.chunk_total = max(1, (self.bytes_total + chunk_size - 1) // chunk_size)
        self._checkpoint_id = make_checkpoint_id(provider, rel_path, self.data_sha, remote_root)
        self._resume: ResumeState | None = None

    @property
    def enabled(self) -> bool:
        return self.progress is not None or self.checkpoint is not None

    def begin(self) -> ResumeState:
        if not self.enabled:
            return ResumeState()
        completed: list[int] = []
        provider_state: dict[str, Any] = {}
        existing = None
        if self.checkpoint is not None:
            existing = self.checkpoint.load(
                provider=self.provider,
                rel_path=self.rel_path,
                remote_root=self.remote_root,
                data_sha=self.data_sha,
            )
        if existing:
            completed = list(existing.get("completed_chunks") or [])
            provider_state = dict(existing.get("provider_state") or {})
            self._checkpoint_id = str(existing.get("checkpoint_id") or self._checkpoint_id)
        start_index = max(completed) + 1 if completed else 0
        if self.progress is not None:
            self.job_id = self.progress.start_job(
                provider=self.provider,
                rel_path=self.rel_path,
                bytes_total=self.bytes_total,
                chunk_total=self.chunk_total,
                checkpoint_id=self._checkpoint_id,
                job_id=self.job_id,
            )
        self._resume = ResumeState(
            checkpoint_id=self._checkpoint_id,
            job_id=self.job_id,
            provider_state=provider_state,
            completed_chunks=completed,
            start_chunk_index=start_index,
        )
        return self._resume

    def after_chunk(self, chunk_index: int, *, provider_state: dict[str, Any]) -> None:
        if not self.enabled or self._resume is None:
            return
        completed = list(self._resume.completed_chunks)
        if chunk_index not in completed:
            completed.append(chunk_index)
        completed.sort()
        self._resume.completed_chunks = completed
        self._resume.provider_state = dict(provider_state)
        bytes_done = min(self.bytes_total, (chunk_index + 1) * self.chunk_size)
        if self.progress is not None and self.job_id:
            try:
                self.progress.on_chunk_done(self.job_id, bytes_done=bytes_done, chunk_index=chunk_index)
            except Exception as exc:
                logger.debug("upload progress update skipped: %s", exc)
        if self.checkpoint is not None:
            try:
                self.checkpoint.save(
                    {
                        "checkpoint_id": self._checkpoint_id,
                        "provider": self.provider,
                        "rel_path": self.rel_path,
                        "remote_root": self.remote_root,
                        "content_sha256": self.data_sha,
                        "bytes_total": self.bytes_total,
                        "chunk_size_bytes": self.chunk_size,
                        "chunk_total": self.chunk_total,
                        "completed_chunks": completed,
                        "provider_state": provider_state,
                    }
                )
            except Exception as exc:
                logger.debug("upload checkpoint save skipped: %s", exc)

    def done(self) -> None:
        if not self.enabled:
            return
        if self.progress is not None and self.job_id:
            try:
                self.progress.finish_job(self.job_id, ok=True)
            except Exception as exc:
                logger.debug("upload progress finish skipped: %s", exc)
        if self.checkpoint is not None and self._checkpoint_id:
            try:
                self.checkpoint.delete(self._checkpoint_id)
            except Exception as exc:
                logger.debug("upload checkpoint delete skipped: %s", exc)

    def failed(self, error: str) -> None:
        if not self.enabled:
            return
        if self.progress is not None and self.job_id:
            try:
                self.progress.finish_job(self.job_id, ok=False, error=error)
            except Exception as exc:
                logger.debug("upload progress fail skipped: %s", exc)


def resolve_upload_config(config: dict[str, Any] | None) -> dict[str, Any]:
    storage = (config or {}).get("storage") or {}
    cloud = storage.get("cloud") or {}
    upload_cfg = cloud.get("upload") or {}
    return {
        "checkpoint_ttl_days": int(upload_cfg.get("checkpoint_ttl_days", 7)),
        "checkpoint_dir": str(upload_cfg.get("checkpoint_dir", "sync/upload_checkpoints")),
        "progress_file": str(upload_cfg.get("progress_file", "sync/upload_progress.json")),
    }


def make_upload_trackers(
    home: Path,
    config: dict[str, Any] | None = None,
) -> tuple[UploadProgressTracker, UploadCheckpointStore]:
    cfg = resolve_upload_config(config)
    progress = UploadProgressTracker(home, progress_file=cfg["progress_file"])
    checkpoint = UploadCheckpointStore(
        home,
        checkpoint_dir=cfg["checkpoint_dir"],
        ttl_days=cfg["checkpoint_ttl_days"],
    )
    return progress, checkpoint


def find_checkpoint_for_rel_path(
    store: UploadCheckpointStore,
    rel_path: str,
    *,
    provider: str | None = None,
) -> dict[str, Any] | None:
    rel = rel_path.strip("/").replace("\\", "/")
    for payload in store.list_checkpoints(provider=provider):
        if str(payload.get("rel_path", "")).strip("/") == rel:
            return payload
    return None
