"""存储层: 标注数据 + 论文语料, 两种后端同一接口。

- local: 标注存 03_annotation/data/store/<annotator>/annotations/<paper>.json, 论文按 manifest 读本地语料树(开发 / 离线用)
- github: 默认一个私有仓库管所有(0925 定): papers/ 放论文, annotations/<annotator>/ 放各人标注。
  服务器用一个 token(Rong 的 fine-grained PAT)读写, 标注人不需要 GitHub 账号, 只能经服务器读写自己那份。
  (也支持每人一个仓库: ANNOTATORS_JSON 里给某人配 "repo", 其标注写进那个仓库的 annotations/。)

标注文件 = 一篇论文 × 一个标注人一份 JSON(records + fields + events), 每次同步一个 commit,
git 历史本身就是第二层审计痕迹。
"""
import base64
import hashlib
import json
import os
import threading
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import httpx

GITHUB_API = os.environ.get("GITHUB_API_URL", "https://api.github.com")


class ConflictError(Exception):
    """存储里的文件已被别处改过(sha 对不上)。由 DocCache 读回合并后重写, 不直接覆盖。"""


# ---------- 标注仓库 ----------

class LocalAnnotationStore:
    def __init__(self, root: Path):
        self.root = root

    def _path(self, annotator: str, name: str) -> Path:
        return self.root / annotator / "annotations" / name

    def read(self, annotator: str, name: str) -> Tuple[Optional[bytes], Optional[str]]:
        p = self._path(annotator, name)
        if not p.exists():
            return None, None
        data = p.read_bytes()
        return data, str(p.stat().st_mtime_ns)

    def write(self, annotator: str, name: str, data: bytes, sha: Optional[str], message: str) -> str:
        p = self._path(annotator, name)
        cur = str(p.stat().st_mtime_ns) if p.exists() else None
        if cur != sha:
            raise ConflictError(name)
        p.parent.mkdir(parents=True, exist_ok=True)
        tmp = p.with_suffix(p.suffix + ".tmp")
        tmp.write_bytes(data)
        tmp.replace(p)
        return str(p.stat().st_mtime_ns)

    def list(self, annotator: str) -> List[str]:
        d = self.root / annotator / "annotations"
        if not d.is_dir():
            return []
        return sorted(f.name for f in d.iterdir() if f.suffix == ".json")


class GitHubAnnotationStore:
    """annotator → (仓库, 目录)。通过 Contents API 读写, 每次写一个 commit。"""

    def __init__(self, token: str, locations: Dict[str, Tuple[str, str]], branch: Optional[str] = None,
                 client: Optional[httpx.Client] = None):
        self.locations = locations
        self.branch = branch
        self.client = client or httpx.Client(base_url=GITHUB_API, timeout=30)
        self.headers = {"Authorization": f"Bearer {token}",
                        "Accept": "application/vnd.github+json",
                        "X-GitHub-Api-Version": "2022-11-28"}

    def _loc(self, annotator: str) -> Tuple[str, str]:
        if annotator not in self.locations:
            raise KeyError(f"no storage location for annotator {annotator}")
        return self.locations[annotator]

    def _url(self, annotator: str, name: str = "") -> str:
        repo, directory = self._loc(annotator)
        return f"/repos/{repo}/contents/{directory}" + (f"/{name}" if name else "")

    def read(self, annotator: str, name: str) -> Tuple[Optional[bytes], Optional[str]]:
        params = {"ref": self.branch} if self.branch else None
        r = self.client.get(self._url(annotator, name),
                            headers=self.headers, params=params)
        if r.status_code == 404:
            return None, None
        r.raise_for_status()
        meta = r.json()
        if meta.get("encoding") == "base64" and meta.get("content"):
            return base64.b64decode(meta["content"]), meta["sha"]
        # >1MB 的文件 content 为空, 走 raw
        raw = self.client.get(self._url(annotator, name),
                              headers={**self.headers, "Accept": "application/vnd.github.raw+json"},
                              params=params)
        raw.raise_for_status()
        return raw.content, meta["sha"]

    def write(self, annotator: str, name: str, data: bytes, sha: Optional[str], message: str) -> str:
        body = {"message": message, "content": base64.b64encode(data).decode()}
        if sha:
            body["sha"] = sha
        if self.branch:
            body["branch"] = self.branch
        url = self._url(annotator, name)
        r = self.client.put(url, headers=self.headers, json=body)
        if r.status_code in (409, 422):
            # sha 过期(别处改过文件, 如另一个服务器实例): 不覆盖, 交给 DocCache 合并后重写
            raise ConflictError(name)
        r.raise_for_status()
        return r.json()["content"]["sha"]

    def list(self, annotator: str) -> List[str]:
        params = {"ref": self.branch} if self.branch else None
        r = self.client.get(self._url(annotator), headers=self.headers, params=params)
        if r.status_code == 404:
            return []
        r.raise_for_status()
        return sorted(e["name"] for e in r.json() if e["type"] == "file" and e["name"].endswith(".json"))


# ---------- 论文语料 ----------

class LocalPapers:
    """manifest 里的路径相对 root。默认 root = 1.Reproducibility/, manifest = 03_annotation/papers/manifest.json。"""

    def __init__(self, root: Path, manifest_path: Path):
        self.root = root
        self.manifest_path = manifest_path

    @property
    def ref(self) -> str:
        return "local"

    def manifest(self) -> dict:
        return json.loads(self.manifest_path.read_text(encoding="utf-8"))

    def text(self, relpath: str) -> str:
        return (self.root / relpath).read_text(encoding="utf-8", errors="ignore")

    def pdf_path(self, relpath: str) -> Path:
        return self.root / relpath

    # ---- 1003: 上传论文 / 论文池(本地: 写在 manifest 旁的 uploads/ 与 pool.json)

    def pool(self) -> Optional[List[str]]:
        p = self.manifest_path.parent / "pool.json"
        return json.loads(p.read_text(encoding="utf-8"))["papers"] if p.exists() else None

    def set_pool(self, papers: List[str], message: str = ""):
        (self.manifest_path.parent / "pool.json").write_text(
            json.dumps({"papers": papers}, ensure_ascii=False, indent=1), encoding="utf-8")

    def add_paper(self, pid: str, files: Dict[str, bytes], entry: dict, message: str = ""):
        """只增不改: pid 已存在就报错(全文是 char_span 坐标系, 冻结)。entry 里的文件字段用 files 的键名。"""
        m = self.manifest()
        if pid in m["papers"]:
            raise FileExistsError(pid)
        d = self.manifest_path.parent / "uploads" / pid
        d.mkdir(parents=True, exist_ok=True)
        for name, data in files.items():
            (d / name).write_bytes(data)
        rel = lambda n: os.path.relpath(d / n, self.root) if n else None  # noqa: E731
        m["papers"][pid] = {**entry, "fulltext": rel(entry["fulltext"]), "pdf": rel(entry.get("pdf")),
                            "layout": rel(entry.get("layout"))}
        self.manifest_path.write_text(json.dumps(m, ensure_ascii=False, indent=2), encoding="utf-8")


class GitHubPapers:
    """仓库里 subdir/(默认 papers/)放 manifest.json, manifest 里的路径相对 subdir。
    启动时把 ref 钉成「最后一次改动 subdir 的 commit」: 本进程内所有读取都来自同一个快照, sha 写进每份标注,
    保证 char_span 坐标系可追溯(md/txt 冻结纪律)。用最后改动 subdir 的 commit 而不是分支头,
    是因为标注和论文同仓库时每次保存都会产生新 commit, 按分支头记会把"论文没变"误报成快照变了。"""

    def __init__(self, token: str, repo: str, ref: str, cache_dir: Path, subdir: str = "papers",
                 client: Optional[httpx.Client] = None):
        self.repo = repo
        self.branch = ref
        self.subdir = subdir.strip("/")
        self.client = client or httpx.Client(base_url=GITHUB_API, timeout=120)
        self.headers = {"Authorization": f"Bearer {token}",
                        "X-GitHub-Api-Version": "2022-11-28"}
        params = {"sha": ref, "per_page": 1}
        if self.subdir:
            params["path"] = self.subdir
        r = self.client.get(f"/repos/{repo}/commits", params=params,
                            headers={**self.headers, "Accept": "application/vnd.github+json"})
        r.raise_for_status()
        commits = r.json()
        if not commits:
            raise RuntimeError(f"no commits touching {self.subdir or '/'} on {repo}@{ref}")
        self._sha = commits[0]["sha"]
        self.cache_dir = cache_dir / self._sha[:12]
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self._manifest = None
        self._pool = False                       # False = 还没读; None = 仓库里没有 pool.json
        self._text_cache: Dict[str, str] = {}
        self._lock = threading.Lock()

    @property
    def ref(self) -> str:
        return self._sha

    def _raw(self, relpath: str) -> bytes:
        full = f"{self.subdir}/{relpath}" if self.subdir else relpath
        r = self.client.get(f"/repos/{self.repo}/contents/{full}",
                            headers={**self.headers, "Accept": "application/vnd.github.raw+json"},
                            params={"ref": self._sha})
        r.raise_for_status()
        return r.content

    def manifest(self) -> dict:
        if self._manifest is None:
            self._manifest = json.loads(self._raw("manifest.json"))
        return self._manifest

    def text(self, relpath: str) -> str:
        with self._lock:
            if relpath not in self._text_cache:
                self._text_cache[relpath] = self._raw(relpath).decode("utf-8", errors="ignore")
            return self._text_cache[relpath]

    def pdf_path(self, relpath: str) -> Path:
        local = self.cache_dir / relpath
        if not local.exists():
            local.parent.mkdir(parents=True, exist_ok=True)
            tmp = local.with_suffix(".part")
            tmp.write_bytes(self._raw(relpath))
            tmp.replace(local)
        return local

    # ---- 1003: 上传论文 / 论文池。一次上传 = 一个 commit(Git Data API), 之后快照 ref 前移到这个 commit:
    # 已有论文的文件不变(只增不改), 所以旧标注的 char_span 仍有效; 各文档的 papers_refs_seen 会记下这次前移。

    def _api(self, method: str, url: str, **kw):
        r = self.client.request(method, f"/repos/{self.repo}{url}",
                                headers={**self.headers, "Accept": "application/vnd.github+json"}, **kw)
        r.raise_for_status()
        return r.json() if r.content else None

    def _head_json(self, relpath: str):
        """分支最新版的某个 papers/ 文件(不是钉住的快照), 改 manifest / pool 时以它为底, 防丢更新。"""
        full = f"{self.subdir}/{relpath}" if self.subdir else relpath
        r = self.client.get(f"/repos/{self.repo}/contents/{full}", params={"ref": self.branch},
                            headers={**self.headers, "Accept": "application/vnd.github.raw+json"})
        if r.status_code == 404:
            return None
        r.raise_for_status()
        return json.loads(r.content)

    def _commit(self, files: Dict[str, bytes], message: str):
        with self._lock:
            blobs = {p: self._api("POST", "/git/blobs", json={"content": base64.b64encode(b).decode(),
                                                               "encoding": "base64"})["sha"]
                     for p, b in files.items()}
            for _ in range(5):                       # 标注同步也在往同一分支提交: 非快进就重来
                head = self._api("GET", f"/git/ref/heads/{self.branch}")["object"]["sha"]
                base_tree = self._api("GET", f"/git/commits/{head}")["tree"]["sha"]
                tree = self._api("POST", "/git/trees", json={"base_tree": base_tree, "tree": [
                    {"path": (f"{self.subdir}/{p}" if self.subdir else p), "mode": "100644", "type": "blob",
                     "sha": s} for p, s in blobs.items()]})["sha"]
                commit = self._api("POST", "/git/commits",
                                   json={"message": message, "tree": tree, "parents": [head]})["sha"]
                r = self.client.patch(f"/repos/{self.repo}/git/refs/heads/{self.branch}", json={"sha": commit},
                                      headers={**self.headers, "Accept": "application/vnd.github+json"})
                if r.status_code == 422:
                    continue
                r.raise_for_status()
                self._sha = commit
                self.cache_dir = self.cache_dir.parent / commit[:12]
                self.cache_dir.mkdir(parents=True, exist_ok=True)
                self._manifest = None
                self._pool = False
                return commit
            raise RuntimeError("could not commit to GitHub (branch kept moving)")

    def pool(self) -> Optional[List[str]]:
        if self._pool is False:
            r = self.client.get(f"/repos/{self.repo}/contents/{self.subdir}/pool.json", params={"ref": self._sha},
                                headers={**self.headers, "Accept": "application/vnd.github.raw+json"})
            self._pool = json.loads(r.content)["papers"] if r.status_code == 200 else None
        return self._pool

    def set_pool(self, papers: List[str], message: str = "papers: update annotator pool"):
        self._commit({"pool.json": json.dumps({"papers": papers}, ensure_ascii=False, indent=1).encode()}, message)

    def add_paper(self, pid: str, files: Dict[str, bytes], entry: dict, message: str = ""):
        m = self._head_json("manifest.json") or {"papers": {}}
        if pid in m["papers"]:
            raise FileExistsError(pid)
        hashes = {n: hashlib.sha256(b).hexdigest() for n, b in files.items()}
        rel = lambda n: f"{pid}/{n}" if n else None  # noqa: E731
        m["papers"][pid] = {**entry, "fulltext": rel(entry["fulltext"]), "pdf": rel(entry.get("pdf")),
                            "layout": rel(entry.get("layout")), "fulltext_sha256": hashes[entry["fulltext"]]}
        out = {f"{pid}/{n}": b for n, b in files.items()}
        out["manifest.json"] = json.dumps(m, ensure_ascii=False, indent=2).encode()
        return self._commit(out, message or f"papers: add {pid}")


# ---------- 写回缓存 ----------

def merge_docs(mine: dict, theirs: dict) -> None:
    """把存储里的版本(theirs)合并进内存版本(mine), 原地修改 mine。
    记录 / 字段按 updated_ts 取较新的一份; 两边都有的保留较新; 只一边有的保留。events、papers_refs_seen 取并集。"""
    def newer(a, b):
        return a if (a or {}).get("updated_ts", 0) >= (b or {}).get("updated_ts", 0) else b

    for uid, rec in (theirs.get("records") or {}).items():
        mine.setdefault("records", {})[uid] = newer(mine["records"].get(uid), rec) if uid in mine["records"] else rec
    for uid, flds in (theirs.get("fields") or {}).items():
        mf = mine.setdefault("fields", {}).setdefault(uid, {})
        for fid, f in flds.items():
            mf[fid] = newer(mf.get(fid), f) if fid in mf else f
    seen = {(e.get("ts"), e.get("kind"), e.get("detail")) for e in mine.get("events", [])}
    extra = [e for e in theirs.get("events", []) if (e.get("ts"), e.get("kind"), e.get("detail")) not in seen]
    if extra:
        mine["events"] = sorted(mine.get("events", []) + extra, key=lambda e: e.get("ts", 0))
    refs = mine.setdefault("papers_refs_seen", [])
    refs += [r for r in theirs.get("papers_refs_seen", []) if r not in refs]
    mine["events"].append({"ts": time.time(), "kind": "merged_concurrent_copy",
                           "detail": f"{len(extra)} events from a concurrently written copy"})


class DocCache:
    """(annotator, paper) → 标注文档, 内存里改, 后台线程每 flush_seconds 秒把脏文档写回存储。
    每次改动都立即在内存生效(界面零延迟); 同步状态通过 status() 暴露给前端的同步指示。"""

    def __init__(self, store, flush_seconds: float = 15.0):
        self.store = store
        self.flush_seconds = flush_seconds
        self.docs: Dict[Tuple[str, str], dict] = {}
        self.shas: Dict[Tuple[str, str], Optional[str]] = {}
        self.dirty: Dict[Tuple[str, str], int] = {}      # key → 脏版本号
        self.pending_msgs: Dict[Tuple[str, str], List[str]] = {}
        self.lock = threading.RLock()
        self.flush_lock = threading.Lock()
        self.last_ok: Dict[str, float] = {}
        self.last_error: Dict[str, Optional[str]] = {}
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None

    @staticmethod
    def path_for(paper: str) -> str:
        return f"{paper}.json"

    def get(self, annotator: str, paper: str, papers_ref: str) -> dict:
        key = (annotator, paper)
        with self.lock:
            if key in self.docs:
                return self.docs[key]
        data, sha = self.store.read(annotator, self.path_for(paper))
        with self.lock:
            if key in self.docs:           # 并发加载, 以先到的为准
                return self.docs[key]
            if data:
                doc = json.loads(data)
            else:
                doc = {"paper": paper, "annotator": annotator, "papers_ref": papers_ref,
                       "created_ts": time.time(), "records": {}, "fields": {}, "events": []}
            self.docs[key] = doc
            self.shas[key] = sha
            return doc

    def mark_dirty(self, annotator: str, paper: str, msg: str):
        key = (annotator, paper)
        with self.lock:
            self.dirty[key] = self.dirty.get(key, 0) + 1
            self.pending_msgs.setdefault(key, []).append(msg)

    def flush(self) -> None:
        with self.flush_lock:
            with self.lock:
                work = [(k, v, json.dumps(self.docs[k], ensure_ascii=False, indent=1).encode("utf-8"),
                         self.pending_msgs.get(k, []))
                        for k, v in self.dirty.items()]
            for key, version, data, msgs in work:
                annotator, paper = key
                n = len(msgs)
                message = f"{annotator} · {paper} · {n} change{'s' if n != 1 else ''}\n\n" + "\n".join(msgs[:50])
                try:
                    try:
                        sha = self.store.write(annotator, self.path_for(paper), data, self.shas.get(key), message)
                    except ConflictError:
                        # 别处写过这份文件: 读回、按字段合并进内存版本, 再写一次(不丢任何一边的改动)
                        remote, rsha = self.store.read(annotator, self.path_for(paper))
                        with self.lock:
                            if remote:
                                merge_docs(self.docs[key], json.loads(remote))
                            self.shas[key] = rsha
                            data = json.dumps(self.docs[key], ensure_ascii=False, indent=1).encode("utf-8")
                        sha = self.store.write(annotator, self.path_for(paper), data, rsha, message + "\n(merged)")
                except Exception as e:                       # 保持脏, 下一轮重试
                    self.last_error[annotator] = f"{type(e).__name__}: {e}"[:300]
                    continue
                with self.lock:
                    self.shas[key] = sha
                    if self.dirty.get(key) == version:       # 期间没有新改动才清脏
                        self.dirty.pop(key, None)
                        self.pending_msgs.pop(key, None)
                    else:
                        self.pending_msgs[key] = self.pending_msgs.get(key, [])[n:]
                    self.last_ok[annotator] = time.time()
                    self.last_error[annotator] = None

    def status(self, annotator: str) -> dict:
        with self.lock:
            pending = sum(1 for (a, _p) in self.dirty if a == annotator)
        return {"pending": pending, "last_ok": self.last_ok.get(annotator),
                "error": self.last_error.get(annotator), "flush_seconds": self.flush_seconds}

    def _loop(self):
        while not self._stop.wait(self.flush_seconds):
            self.flush()

    def start(self):
        if self._thread is None:
            self._thread = threading.Thread(target=self._loop, name="doc-flusher", daemon=True)
            self._thread.start()

    def stop(self):
        self._stop.set()
        self.flush()


def env_flag(name: str, default: bool) -> bool:
    v = os.environ.get(name)
    if v is None:
        return default
    return v.strip().lower() in ("1", "true", "yes", "on")