"""Export normalized Ministry fragments for the OpenSearch Bulk API."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sqlite3
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Iterator

APP_VERSION = "0.5.3"
LANGUAGES = {"ru", "ky"}
DEFAULT_MAPPING = Path(__file__).with_name("minjust-fragments-index.json")


def read_json(path: Path):
    def reject_constant(value: str):
        raise ValueError(f"Invalid JSON constant: {value}")

    for attempt in range(3):
        try:
            with path.open(encoding="utf-8") as source:
                return json.load(source, parse_constant=reject_constant)
        except (OSError, UnicodeError, json.JSONDecodeError) as error:
            if attempt == 2:
                raise ValueError(f"Cannot read JSON {path}: {error}") from error
            time.sleep(0.1)


def localized(value: object, language: str):
    return value.get(language) if isinstance(value, dict) else None


def paths(document: dict, field: str, language: str) -> list[str]:
    return [" > ".join(item[language]) for item in document.get(field, []) if item.get(language)]


def search_document(document: dict, fragment: dict, expected: tuple[str, str, str, int]) -> dict:
    document_code, edition_code, language, position = expected
    fragment_id = f"document:{document_code}:edition:{edition_code}:lang:{language}:fragment:{position}"
    required = ("id", "document_code", "edition_code", "language", "position", "type", "text", "text_sha256", "source_path", "source_sha256")
    if not isinstance(fragment, dict) or any(key not in fragment for key in required):
        raise ValueError(f"Incomplete fragment {fragment_id}")
    if (fragment["document_code"], fragment["edition_code"], fragment["language"], fragment["position"], fragment["id"]) != (*expected, fragment_id):
        raise ValueError(f"Fragment identity does not match its path: {fragment_id}")
    if language not in LANGUAGES or not isinstance(fragment["text"], str) or not fragment["text"]:
        raise ValueError(f"Invalid fragment content: {fragment_id}")
    for key in ("text_sha256", "source_sha256"):
        value = fragment[key]
        if not isinstance(value, str) or len(value) != 64 or any(char not in "0123456789abcdef" for char in value):
            raise ValueError(f"Invalid {key}: {fragment_id}")
    if hashlib.sha256(fragment["text"].encode()).hexdigest() != fragment["text_sha256"]:
        raise ValueError(f"Text checksum mismatch: {fragment_id}")

    dates = document.get("dates") or {}
    result = {
        "schema_version": document["schema_version"],
        "document_code": document_code,
        "edition_code": edition_code,
        "language": language,
        "position": position,
        "fragment_type": fragment["type"],
        f"text_{language}": fragment["text"],
        "document_name_ru": localized(document.get("name"), "ru"),
        "document_name_ky": localized(document.get("name"), "ky"),
        "document_type_ru": localized(document.get("type"), "ru"),
        "document_type_ky": localized(document.get("type"), "ky"),
        "status_ru": localized(document.get("status"), "ru"),
        "status_ky": localized(document.get("status"), "ky"),
        "number": document.get("number"),
        "date_adopted": dates.get("DateAdopted"),
        "authority_paths_ru": paths(document, "authority_paths", "ru"),
        "authority_paths_ky": paths(document, "authority_paths", "ky"),
        "source_path": fragment["source_path"],
        "source_sha256": fragment["source_sha256"],
        "text_sha256": fragment["text_sha256"],
    }
    return {key: value for key, value in result.items() if value is not None}


def document_codes(input_root: Path, limit: int | None = None, start_at: str | None = None) -> list[str]:
    connection = sqlite3.connect(f"file:{input_root / 'manifest.sqlite3'}?mode=ro", uri=True)
    try:
        query = "SELECT code FROM documents WHERE state = 'success' ORDER BY CAST(code AS INTEGER), code"
        parameters: list[int] = []
        if limit is not None:
            query += " LIMIT ?"
            parameters.append(limit)
        codes = [str(code) for (code,) in connection.execute(query, parameters)]
    finally:
        connection.close()
    if start_at is None:
        return codes
    try:
        return codes[codes.index(start_at):]
    except ValueError as error:
        raise ValueError(f"Resume document not found in selected range: {start_at}") from error


def bulk_pairs(
    input_root: Path,
    index: str,
    limit: int | None = None,
    start_at: str | None = None,
) -> Iterator[tuple[str, bytes]]:
    document_root = input_root / "documents"
    if not document_root.is_dir():
        raise FileNotFoundError(f"Document directory not found: {document_root}")
    for code in document_codes(input_root, limit, start_at):
        directory = document_root / code
        document = read_json(directory / "document.json")
        if document.get("source_code") != directory.name:
            raise ValueError(f"Document identity does not match its path: {directory}")
        for fragment_path in sorted(directory.glob("editions/*/*/fragments.json")):
            edition_code, language = fragment_path.parts[-3:-1]
            values = read_json(fragment_path)
            if not isinstance(values, list):
                raise ValueError(f"Fragments must be a list: {fragment_path}")
            for position, fragment in enumerate(values, 1):
                source = search_document(document, fragment, (directory.name, edition_code, language, position))
                action = json.dumps({"index": {"_index": index, "_id": fragment["id"]}}, ensure_ascii=False, allow_nan=False)
                body = json.dumps(source, ensure_ascii=False, allow_nan=False)
                yield directory.name, f"{action}\n{body}\n".encode()


def document_count(input_root: Path, limit: int | None, start_at: str | None = None) -> int:
    return len(document_codes(input_root, limit, start_at))


def bulk_batches(pairs: Iterator[tuple[str, bytes]], maximum_bytes: int) -> Iterator[tuple[str, bytes]]:
    batch = bytearray()
    last_code = ""
    for code, pair in pairs:
        if len(pair) > maximum_bytes:
            raise ValueError(f"One Bulk pair exceeds the {maximum_bytes}-byte batch limit")
        if batch and len(batch) + len(pair) > maximum_bytes:
            yield last_code, bytes(batch)
            batch.clear()
        last_code = code
        batch.extend(pair)
    if batch:
        yield last_code, bytes(batch)


def export_bulk(input_root: Path, output: Path, index: str, limit: int | None = None) -> tuple[int, int]:
    source = input_root.resolve()
    destination = output.resolve()
    if destination == source or destination.is_relative_to(source):
        raise ValueError("--output must not be inside --input")
    output.parent.mkdir(parents=True, exist_ok=True)
    fragments = 0
    with tempfile.NamedTemporaryFile("wb", dir=output.parent, delete=False) as target:
        temporary = Path(target.name)
        try:
            for code, pair in bulk_pairs(input_root, index, limit):
                target.write(pair)
                fragments += 1
            target.flush()
            os.fsync(target.fileno())
            os.replace(temporary, output)
        except BaseException:
            temporary.unlink(missing_ok=True)
            raise
    return document_count(input_root, limit), fragments


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_json_atomic(path: Path, value: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("wb", dir=path.parent, delete=False) as target:
        temporary = Path(target.name)
        try:
            target.write((json.dumps(value, ensure_ascii=False, indent=2) + "\n").encode())
            target.flush()
            os.fsync(target.fileno())
            os.replace(temporary, path)
        except BaseException:
            temporary.unlink(missing_ok=True)
            raise


def request_json(
    url: str,
    method: str,
    body: bytes | None,
    content_type: str,
    attempts: int = 5,
    retry_invalid_json: bool = False,
) -> dict:
    request = urllib.request.Request(url, data=body, method=method, headers={"Content-Type": content_type})
    for attempt in range(attempts):
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                return json.load(response)
        except (UnicodeError, json.JSONDecodeError) as error:
            if not retry_invalid_json or attempt == attempts - 1:
                raise RuntimeError(f"{method} {url} returned invalid JSON: {error}") from error
            delay = 2**attempt
        except urllib.error.HTTPError as error:
            response_body = error.read(500).decode("utf-8", errors="replace")
            error.close()
            retryable = error.code == 429 or 500 <= error.code < 600
            if not retryable or attempt == attempts - 1:
                raise RuntimeError(
                    f"{method} {url} failed with HTTP {error.code}: {response_body}"
                ) from error
            retry_after = error.headers.get("Retry-After")
            delay = float(retry_after) if retry_after and retry_after.isdigit() else 2**attempt
        except (urllib.error.URLError, TimeoutError, ConnectionResetError) as error:
            if attempt == attempts - 1:
                raise RuntimeError(f"{method} {url} failed after {attempts} attempts: {error}") from error
            delay = 2**attempt
        time.sleep(delay)
    raise AssertionError("unreachable")


def opensearch_identity(base: str, index: str, index_url: str) -> tuple[str, str]:
    cluster = request_json(f"{base}/", "GET", None, "application/json")
    definition = request_json(index_url, "GET", None, "application/json")
    try:
        return cluster["cluster_uuid"], definition[index]["settings"]["index"]["uuid"]
    except (KeyError, TypeError) as error:
        raise RuntimeError("OpenSearch identity response is incomplete") from error


def checkpoint_state(
    input_root: Path,
    index: str,
    checkpoint: Path,
    resume: bool,
    url: str,
    cluster_uuid: str,
    index_uuid: str,
    limit: int | None,
) -> tuple[dict, str | None]:
    source = input_root.resolve()
    destination = checkpoint.resolve()
    if destination == source or destination.is_relative_to(source):
        raise ValueError("--checkpoint must not be inside --input")
    manifest_sha256 = file_sha256(input_root / "manifest.sqlite3")
    if not resume:
        return {
            "schema_version": 1,
            "url": url,
            "cluster_uuid": cluster_uuid,
            "index": index,
            "index_uuid": index_uuid,
            "input": str(source),
            "manifest_sha256": manifest_sha256,
            "limit": limit,
            "last_document_code": None,
            "complete": False,
        }, None

    state = read_json(checkpoint)
    expected = {
        "schema_version",
        "url",
        "cluster_uuid",
        "index",
        "index_uuid",
        "input",
        "manifest_sha256",
        "limit",
        "last_document_code",
        "complete",
    }
    if not isinstance(state, dict) or set(state) != expected:
        raise ValueError(f"Invalid checkpoint: {checkpoint}")
    if (
        state["schema_version"] != 1
        or state["url"] != url
        or state["cluster_uuid"] != cluster_uuid
        or state["index"] != index
        or state["index_uuid"] != index_uuid
        or state["input"] != str(source)
    ):
        raise ValueError(f"Checkpoint does not match this load: {checkpoint}")
    if state["limit"] != limit:
        raise ValueError(f"Checkpoint limit does not match --limit: {checkpoint}")
    if state["manifest_sha256"] != manifest_sha256:
        raise ValueError("Normalized manifest changed; create a new versioned index")
    if state["complete"] is not False:
        raise ValueError(f"Checkpoint is already complete: {checkpoint}")
    start_at = state["last_document_code"]
    if start_at is not None and not isinstance(start_at, str):
        raise ValueError(f"Invalid checkpoint document code: {checkpoint}")
    return state, start_at


def load_bulk(
    input_root: Path,
    url: str,
    index: str,
    mapping: Path = DEFAULT_MAPPING,
    maximum_bytes: int = 25 * 1024 * 1024,
    limit: int | None = None,
    resume: bool = False,
    checkpoint: Path = Path("data/opensearch/minjust-fragments.checkpoint.json"),
    alias: str | None = None,
) -> tuple[int, int]:
    base = url.rstrip("/")
    if alias is not None and alias == index:
        raise ValueError("--alias must differ from --index")
    index_url = f"{base}/{urllib.parse.quote(index, safe='')}"
    if resume:
        cluster_uuid, index_uuid = opensearch_identity(base, index, index_url)
    else:
        request_json(index_url, "PUT", mapping.read_bytes(), "application/json")
        cluster_uuid, index_uuid = opensearch_identity(base, index, index_url)
    state, start_at = checkpoint_state(
        input_root,
        index,
        checkpoint,
        resume,
        base,
        cluster_uuid,
        index_uuid,
        limit,
    )
    documents = document_count(input_root, limit, start_at)
    if not resume:
        write_json_atomic(checkpoint, state)
    fragments = 0
    for last_code, batch in bulk_batches(bulk_pairs(input_root, index, limit, start_at), maximum_bytes):
        result = request_json(
            f"{base}/_bulk",
            "POST",
            batch,
            "application/x-ndjson",
            retry_invalid_json=True,
        )
        expected = batch.count(b"\n") // 2
        items = result.get("items", [])
        if result.get("errors"):
            failures = [item.get("index", {}) for item in items if item.get("index", {}).get("error")]
            details = "; ".join(
                f"{item.get('_id', '<unknown>')}: {item['error'].get('type', 'error')}: "
                f"{item['error'].get('reason', '<no reason>')}"
                for item in failures[:5]
            )
            raise RuntimeError(
                f"OpenSearch Bulk API failed for {len(failures)} item(s); "
                f"resume from {checkpoint}: {details}"
            )
        if len(items) != expected:
            raise RuntimeError(
                f"OpenSearch Bulk API returned {len(items)} of {expected} item result(s); "
                f"resume from {checkpoint}"
            )
        fragments += len(items)
        state["last_document_code"] = last_code
        write_json_atomic(checkpoint, state)
        print(f"checkpoint={last_code} fragments={fragments}", flush=True)
    if alias:
        switch_alias(base, index, alias)
    state["complete"] = True
    write_json_atomic(checkpoint, state)
    return documents, fragments


def switch_alias(base: str, index: str, alias: str) -> None:
    if not alias or alias == index:
        raise ValueError("--alias must differ from --index")
    result = request_json(
        f"{base.rstrip('/')}/_aliases",
        "POST",
        json.dumps({
            "actions": [
                {"remove": {"index": "*", "alias": alias, "must_exist": False}},
                {"add": {"index": index, "alias": alias}},
            ]
        }).encode(),
        "application/json",
    )
    if result.get("acknowledged") is not True:
        raise RuntimeError(f"OpenSearch did not acknowledge alias switch: {alias}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=Path("data/minjust-normalized"))
    parser.add_argument("--output", type=Path, default=Path("data/opensearch/minjust-fragments.ndjson"))
    parser.add_argument("--index", default="akyldash-fragments-v1")
    parser.add_argument("--limit", type=int)
    parser.add_argument("--url", help="create the index and stream bounded Bulk requests instead of writing a file")
    parser.add_argument("--alias", help="atomically point this alias at --index after a successful load")
    parser.add_argument("--mapping", type=Path, default=DEFAULT_MAPPING)
    parser.add_argument("--batch-mb", type=int, default=25)
    parser.add_argument("--resume", action="store_true", help="load into an existing index")
    parser.add_argument("--checkpoint", type=Path, help="persistent resume checkpoint path")
    parser.add_argument("--version", action="version", version=APP_VERSION)
    arguments = parser.parse_args()
    if arguments.limit is not None and arguments.limit <= 0:
        raise SystemExit("--limit must be greater than zero")
    if arguments.batch_mb <= 0:
        raise SystemExit("--batch-mb must be greater than zero")
    if arguments.resume and not arguments.url:
        raise SystemExit("--resume requires --url")
    if arguments.checkpoint and not arguments.url:
        raise SystemExit("--checkpoint requires --url")
    if arguments.alias and not arguments.url:
        raise SystemExit("--alias requires --url")
    if arguments.url:
        checkpoint = arguments.checkpoint or Path("data/opensearch") / f"{arguments.index}.checkpoint.json"
        documents, fragments = load_bulk(
            arguments.input,
            arguments.url,
            arguments.index,
            arguments.mapping,
            arguments.batch_mb * 1024 * 1024,
            arguments.limit,
            arguments.resume,
            checkpoint,
            arguments.alias,
        )
        destination = arguments.url
    else:
        documents, fragments = export_bulk(arguments.input, arguments.output, arguments.index, arguments.limit)
        destination = arguments.output
    print(f"documents={documents} fragments={fragments} destination={destination}\nAkyldash Backend v{APP_VERSION} · Frontend — not created")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())