"""Local Word COM tools for refreshable Zotero citations. No cloud credentials."""
import contextlib
import datetime
import hashlib
import json
import pathlib
import re
import shutil
import threading
import time
import urllib.request
import uuid
import xml.etree.ElementTree as ET

import pythoncom
import pywintypes
import win32com.client
import win32gui

LOCK = threading.RLock()
MARKER = re.compile(r"\[@([A-Z0-9]{8}(?:\s*,\s*@[A-Z0-9]{8})*)\]")
API = "http://127.0.0.1:23119/api/users/0"


def _json(url):
    # Loopback only. Never forwards credentials or library content to another host.
    with urllib.request.urlopen(url, timeout=20) as response:
        return json.load(response)


@contextlib.contextmanager
def _document(filename, write=False):
    path = pathlib.Path(filename).expanduser().resolve(strict=True)
    if path.suffix.lower() != ".docx":
        raise ValueError("Provide an existing .docx file with an absolute path.")
    with LOCK:
        pythoncom.CoInitialize()
        app = doc = None
        created_app = opened_doc = False
        try:
            try:
                app = win32com.client.GetActiveObject("Word.Application")
            except pythoncom.com_error:
                app = win32com.client.DispatchEx("Word.Application")
                app.Visible = True
                created_app = True
            # Generated dispatch wrappers keep Word methods distinct from COM properties.
            app = win32com.client.gencache.EnsureDispatch(app)
            for candidate in app.Documents:
                if str(candidate.FullName).casefold() == str(path).casefold():
                    doc = candidate
                    break
            if doc is None:
                doc = app.Documents.Open(str(path), ReadOnly=False, AddToRecentFiles=False)
                opened_doc = True
            if created_app:
                app.Activate()
                time.sleep(0.5)
                # Office registers its documents for other apps after losing focus.
                candidates = []
                win32gui.EnumWindows(lambda h, _: candidates.append(h)
                                     if win32gui.IsWindowVisible(h) and
                                     "Zotero" in win32gui.GetWindowText(h) else None, None)
                if candidates:
                    try:
                        win32gui.SetForegroundWindow(candidates[0])
                    except (pythoncom.com_error, pywintypes.error):
                        pass
                time.sleep(0.5)
            if write:
                if doc.ReadOnly:
                    raise ValueError("Word opened the document read-only; no changes made.")
                backup_dir = path.parent / ".word-zotero-backups"
                backup_dir.mkdir(exist_ok=True)
                stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S-%f")
                shutil.copy2(path, backup_dir / (path.stem + "-" + stamp + ".docx"))
            doc.Activate()
            yield app, doc
            if write:
                doc.Save()
        finally:
            try:
                if opened_doc and doc is not None:
                    doc.Close(0)
            finally:
                doc = None
                try:
                    if created_app and app is not None:
                        app.Quit(0)
                finally:
                    app = None
                    pythoncom.CoUninitialize()


def _fields(doc):
    result = []
    for field in doc.Fields:
        code = str(field.Code.Text).strip()
        if "ZOTERO_ITEM CSL_CITATION" in code:
            payload = json.loads(code.split("CSL_CITATION", 1)[1].strip())
            result.append({"kind": "citation", "text": str(field.Result.Text),
                           "citation_id": payload.get("citationID"),
                           "uris": [uri for item in payload.get("citationItems", [])
                                    for uri in item.get("uris", [])]})
        elif "ZOTERO_BIBL" in code:
            result.append({"kind": "bibliography", "text": str(field.Result.Text)})
    return result


def _preferences(doc, style_id=None):
    chunks = []
    i = 1
    while True:
        try:
            chunks.append(str(doc.CustomDocumentProperties.Item(f"ZOTERO_PREF_{i}").Value))
        except pythoncom.com_error:
            break
        i += 1
    old = "".join(chunks)
    if old and not style_id:
        return old
    if style_id:
        style_id = style_id.removeprefix("http://www.zotero.org/styles/")
        if not re.fullmatch(r"[a-z0-9-]+", style_id):
            raise ValueError("Provide a Zotero CSL style name, e.g. ieee or apa.")
        if not (pathlib.Path.home() / "Zotero/styles" / (style_id + ".csl")).exists():
            raise ValueError("Citation style is not installed in local Zotero: " + style_id)
    style_url = "http://www.zotero.org/styles/" + (style_id or "ieee")
    if old.startswith("{"):
        data = json.loads(old)
        data["style"]["styleID"] = style_url
        data["style"]["bibliographyStyleHasBeenSet"] = False
        value = json.dumps(data)
    else:
        root = ET.fromstring(old) if old else ET.fromstring(
            '<data data-version="3" zotero-version="10.0.6">'
            f'<session id="{uuid.uuid4().hex}"/>'
            '<style locale="en-US" hasBibliography="1" bibliographyStyleHasBeenSet="0"/>'
            '<prefs><pref name="fieldType" value="Field"/>'
            '<pref name="noteType" value="0"/>'
            '<pref name="automaticJournalAbbreviations" value="false"/></prefs></data>')
        root.find("style").set("id", style_url)
        root.find("style").set("bibliographyStyleHasBeenSet", "0")
        value = ET.tostring(root, encoding="unicode")
    for prop in list(doc.CustomDocumentProperties):
        if re.fullmatch(r"ZOTERO_PREF_\d+", str(prop.Name)):
            prop.Delete()
    for i, start in enumerate(range(0, len(value), 255), 1):
        doc.CustomDocumentProperties.Add(f"ZOTERO_PREF_{i}", False, 4, value[start:start + 255])
    return value


def _citation(keys):
    if not keys or len(keys) != len(set(keys)):
        raise ValueError("Provide a nonempty list of unique Zotero item keys.")
    items = []
    for key in keys:
        if not re.fullmatch(r"[A-Z0-9]{8}", key):
            raise ValueError("Invalid Zotero item key: " + key)
        csl = _json(f"{API}/items/{key}?format=csljson")[0]
        uri = csl["id"]
        numeric_id = int(hashlib.sha256(uri.encode()).hexdigest()[:8], 16)
        csl["id"] = numeric_id
        items.append({"id": numeric_id, "uris": [uri], "itemData": csl})
    return {"citationID": uuid.uuid4().hex[:12], "citationItems": items,
            "properties": {"formattedCitation": "[?]", "plainCitation": "[?]", "noteIndex": 0},
            "schema": "https://github.com/citation-style-language/schema/raw/master/csl-citation.json"}


def _add_field(doc, target, instruction, display):
    if "ZOTERO_ITEM" in instruction and target.Start > 0:
        previous = doc.Range(target.Start - 1, target.Start)
        if previous.Text == " ":
            previous.Text = "\u00a0"
        elif not previous.Text.isspace():
            start = target.Start
            target.InsertBefore("\u00a0")
            target.Start = start + 1
    # Zotero itself uses a QUOTE field and then replaces its code (Word range stability).
    # A nonempty QUOTE is essential: an empty one has no result separator,
    # so assigning Result.Text writes outside the field instead.
    field = doc.Fields.Add(target, 35, '"ZOTERO_PLACEHOLDER"', False)
    field.Result.Text = display
    field.Code.Text = " " + instruction + " "
    field.ShowCodes = False
    return field


def _find(doc, text):
    if not text:
        raise ValueError("An exact nonempty anchor is required.")
    target = doc.Content.Duplicate
    if not target.Find.Execute(FindText=text, MatchCase=True, MatchWildcards=False, Wrap=0):
        raise ValueError("Anchor was not found: " + text)
    second = doc.Range(target.End, doc.Content.End)
    if second.Find.Execute(FindText=text, MatchCase=True, MatchWildcards=False, Wrap=0):
        raise ValueError("Anchor is ambiguous; provide a longer unique sentence.")
    return target


def _outside_fields(doc, target):
    for field in doc.Fields:
        if target.Start < field.Result.End + 1 and target.End > field.Code.Start - 1:
            raise ValueError("The requested text overlaps a Word field; edit the citation through Zotero.")


def _refresh(app, doc, timeout=45):
    if not any(str(a.Name).lower() == "zotero.dotm" and a.Installed for a in app.AddIns):
        raise RuntimeError("Zotero.dotm is not loaded in Word.")
    doc.Activate()
    app.Run("Zotero.ZoteroRefresh")
    started = time.monotonic()
    previous = None
    stable = 0
    while time.monotonic() - started < timeout:
        pythoncom.PumpWaitingMessages()
        try:
            state = _fields(doc)
            text = json.dumps(state, ensure_ascii=False)
            complete = state and all(x["text"].strip() and "[?]" not in x["text"]
                                     and "{Bibliography}" not in x["text"] for x in state)
            stable = stable + 1 if complete and text == previous else 0
            previous = text
            if stable >= 3:
                doc.Save()
                return {"refreshed": True, "fields": state}
        except pythoncom.com_error:
            stable = 0
        time.sleep(0.5)
    preview = [{"kind": x["kind"], "text": x["text"][:160]} for x in _fields(doc)]
    raise TimeoutError("Zotero refresh did not complete. A backup exists. State: " + json.dumps(preview))


def word_live_read(filename: str) -> dict:
    """Read actual Word paragraphs and live Zotero fields; works for an open .docx too."""
    with _document(filename) as (app, doc):
        return {"filename": str(doc.FullName), "word_version": str(app.Version),
                "paragraphs": [{"index": i, "text": str(p.Range.Text).rstrip("\r\x07")}
                               for i, p in enumerate(doc.Paragraphs, 1)], "fields": _fields(doc)}


def word_live_replace(filename: str, old_text: str, new_text: str) -> dict:
    """Replace one unique exact passage using Word COM, preserving citation fields; creates backup."""
    with _document(filename, write=True) as (_, doc):
        target = _find(doc, old_text)
        _outside_fields(doc, target)
        target.Text = new_text
        return {"replaced": True, "filename": str(doc.FullName)}


def word_zotero_insert_citation(filename: str, anchor_text: str, item_keys: list[str],
                                position: str = "after") -> dict:
    """Insert refreshable library-linked Zotero citation before/after a unique exact passage, then refresh."""
    if position not in ("before", "after"):
        raise ValueError("position must be before or after")
    payload = _citation(item_keys)
    with _document(filename, write=True) as (app, doc):
        target = _find(doc, anchor_text)
        target.Collapse(0 if position == "after" else 1)
        _outside_fields(doc, target)
        _preferences(doc)
        _add_field(doc, target, "ADDIN ZOTERO_ITEM CSL_CITATION " + json.dumps(payload), "[?]")
        return _refresh(app, doc)


def word_zotero_apply_markers(filename: str, add_bibliography: bool = True) -> dict:
    """Convert [@KEY] or [@KEY1, @KEY2] markers into live Zotero citations and refresh with Zotero."""
    with _document(filename, write=True) as (app, doc):
        markers = list(MARKER.finditer(str(doc.Content.Text)))
        prepared = [(m.group(0), _citation(re.findall(r"[A-Z0-9]{8}", m.group(1)))) for m in markers]
        _preferences(doc)
        for marker, payload in prepared:
            target = doc.Content.Duplicate
            if not target.Find.Execute(FindText=marker, MatchWildcards=False, Wrap=0):
                raise ValueError("Citation marker disappeared during editing.")
            _outside_fields(doc, target)
            _add_field(doc, target, "ADDIN ZOTERO_ITEM CSL_CITATION " + json.dumps(payload), "[?]")
        if add_bibliography and not any(x["kind"] == "bibliography" for x in _fields(doc)):
            target = doc.Range(doc.Content.End - 1, doc.Content.End - 1)
            _add_field(doc, target, 'ADDIN ZOTERO_BIBL {"uncited":[],"omitted":[],"custom":[]} CSL_BIBLIOGRAPHY', "{Bibliography}")
        result = _refresh(app, doc)
        result["markers_converted"] = len(prepared)
        return result


def word_zotero_refresh(filename: str) -> dict:
    """Run the installed Zotero Word add-in Refresh and save after fields have settled."""
    with _document(filename, write=True) as (app, doc):
        return _refresh(app, doc)


def word_zotero_set_style(filename: str, style_id: str) -> dict:
    """Set an installed Zotero CSL style (ieee, apa, etc.) and refresh citations and bibliography."""
    with _document(filename, write=True) as (app, doc):
        _preferences(doc, style_id)
        return _refresh(app, doc)


def register(mcp):
    for function in (word_live_read, word_live_replace, word_zotero_insert_citation,
                     word_zotero_apply_markers, word_zotero_refresh, word_zotero_set_style):
        mcp.tool(annotations={"readOnlyHint": function is word_live_read})(function)