"""Bounded-memory validation and import of large graph-data JSON files.

The browser never needs to parse these files. SQLite is a temporary index,
not another graph database: it validates cross-record references and supplies
bounded UNWIND batches. Only the completed FalkorDB graph is published.
"""
from __future__ import annotations

import hashlib
import json
import math
import sqlite3
from collections.abc import Mapping
from contextlib import nullcontext
from functools import lru_cache
from pathlib import Path
from typing import BinaryIO, Callable

import ijson
from ijson.common import ObjectBuilder

from app.project_lifecycle import (
    GraphImportError, ProjectValidationError, SpatialMapValidation,
    _falkor_client, _graph_safe_properties, _node_labels,
    normalize_graph_node, normalize_graph_relation, normalize_provision_header,
)

Progress = Callable[[str, int, int], None]
MAX_STRING_BYTES = 4 * 1024 * 1024
MAX_RECORD_CHARS = 8 * 1024 * 1024


class CheckedReader:
    """Reject oversized tokens/deep nesting before the parser buffers them."""

    def __init__(self, stream):
        self.stream = stream
        self.sha256 = hashlib.sha256()
        self.count = 0
        self.in_string = self.escaped = False
        self.token_size = self.depth = 0

    def read(self, size=-1):
        block = self.stream.read(65536 if size < 0 else min(size, 65536))
        self.sha256.update(block)
        self.count += len(block)
        for char in block:
            if self.in_string:
                self.token_size += 1
                if self.token_size > MAX_STRING_BYTES:
                    raise ValueError("单个 JSON 字符串超过 4 MiB，请将大正文/附件拆分为记录或外部资源")
                if self.escaped:
                    self.escaped = False
                elif char == 92:
                    self.escaped = True
                elif char == 34:
                    self.in_string = False
                    self.token_size = 0
            elif char == 34:
                self.in_string = True
                self.token_size = 0
            elif char in (123, 91):
                self.depth += 1
                self.token_size = 0
                if self.depth > 64:
                    raise ValueError("JSON 嵌套不能超过 64 层")
            elif char in (125, 93):
                self.depth -= 1
                self.token_size = 0
            elif char in (9, 10, 13, 32, 44, 58):
                self.token_size = 0
            else:
                self.token_size += 1
                if self.token_size > 1024:
                    raise ValueError("JSON 数值或标记过长")
        return block


def graph_records(reader):
    """One pass, strict root/array shape and duplicate-key detection."""
    builder = None
    record_kind = ""
    record_depth = 0
    record_chars = 0
    maps: list[set[str] | None] = []
    root_fields: set[str] = set()
    first = True
    for prefix, event, value in ijson.parse(reader, use_float=True):
        if first:
            first = False
            if (prefix, event) != ("", "start_map"):
                raise ValueError("图谱数据 JSON 顶层必须是对象")
        if event == "map_key":
            keys = maps[-1]
            if value in keys:
                raise ValueError(f"JSON 对象中存在重复字段：{str(value)[:120]}")
            keys.add(value)
            if len(keys) > 100000:
                raise ValueError("单个 JSON 对象字段过多")
            if prefix == "":
                root_fields.add(value)
        elif event in ("start_map", "start_array"):
            maps.append(set() if event == "start_map" else None)
        elif event in ("end_map", "end_array"):
            maps.pop()
        if prefix in ("nodes", "relations") and event not in ("start_array", "end_array", "map_key"):
            raise ValueError(f"graph_data.{prefix} 必须是数组")
        if builder is None and prefix in ("nodes.item", "relations.item"):
            if event != "start_map":
                raise ValueError(f"{prefix} 必须是 JSON 对象")
            builder, record_kind, record_depth, record_chars = ObjectBuilder(), prefix.split(".")[0], 0, 0
        if builder is not None:
            record_chars += len(value) if isinstance(value, str) else 8
            if record_chars > MAX_RECORD_CHARS:
                raise ValueError("单条节点或关系超过 8 MiB，请拆分大附件")
            builder.event(event, value)
            if event in ("start_map", "start_array"):
                record_depth += 1
            elif event in ("end_map", "end_array"):
                record_depth -= 1
            if record_depth == 0:
                yield record_kind, builder.value
                builder = None
    if "nodes" not in root_fields:
        raise ValueError("图谱数据缺少 nodes 数组；请上传 data JSON，而不是 Schema 或 bundle 文件")


class NodeTypes(Mapping):
    """A bounded LRU over SQLite, rather than an all-nodes Python dictionary."""

    def __init__(self, connection):
        self.connection = connection
        self.lookup = lru_cache(maxsize=4096)(self._lookup)

    def _lookup(self, key):
        row = self.connection.execute("SELECT type FROM nodes WHERE id=?", (key,)).fetchone()
        if row is None:
            raise KeyError(key)
        return row[0]

    def __getitem__(self, key):
        return self.lookup(key)

    def __iter__(self):
        return (row[0] for row in self.connection.execute("SELECT id FROM nodes"))

    def __len__(self):
        return self.connection.execute("SELECT count(*) FROM nodes").fetchone()[0]


def _dumps(value):
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False)


def validate_graph_file(source: Path | BinaryIO, staging: Path, body: dict,
                        progress: Progress = lambda *_: None) -> dict:
    """Fully validate before any PostgreSQL project/FalkorDB write."""
    payload = normalize_provision_header(body)
    if staging.exists():
        raise FileExistsError("校验暂存文件已存在")
    connection = sqlite3.connect(staging)
    node_types = NodeTypes(connection)
    spatial = SpatialMapValidation(payload["spatial_map"])
    counts = {"nodes": 0, "relations": 0}
    try:
        connection.execute("PRAGMA cache_size=-4096")
        connection.execute("PRAGMA temp_store=FILE")
        connection.executescript("""
            CREATE TABLE nodes(id TEXT PRIMARY KEY, type TEXT NOT NULL, labels TEXT NOT NULL, payload TEXT NOT NULL);
            CREATE TABLE relations(seq INTEGER PRIMARY KEY, type TEXT, source_type TEXT, target_type TEXT, payload TEXT NOT NULL);
        """)
        # Multipart uploads are already spooled to disk; reuse their handle
        # instead of making another full-size string or file copy.
        with (source.open("rb") if isinstance(source, Path) else nullcontext(source)) as stream:
            stream.seek(0)
            reader = CheckedReader(stream)
            for kind, raw in graph_records(reader):
                errors: list[str] = []
                if kind == "nodes":
                    node = normalize_graph_node(raw, counts[kind], payload["schema"]["entity_types"], node_types, errors)
                    if node is not None:
                        spatial.add(node, errors)
                    if errors:
                        raise ProjectValidationError(errors[:100])
                    connection.execute(
                        "INSERT INTO nodes VALUES(?,?,?,?)",
                        (node["id"], node["type"], _dumps(_node_labels(node)), _dumps(node)),
                    )
                else:
                    # Relations may appear before nodes in valid JSON.
                    connection.execute("INSERT INTO relations(seq,payload) VALUES(?,?)", (counts[kind], _dumps(raw)))
                counts[kind] += 1
                if (counts["nodes"] + counts["relations"]) % 1000 == 0:
                    connection.commit()
                    progress("校验节点" if kind == "nodes" else "读取关系", counts["nodes"], counts["relations"])
        if not counts["nodes"]:
            raise ProjectValidationError(["graph_data.nodes 至少需要一个节点，FalkorDB 不支持持久化空图"])
        errors = []
        spatial.finish(errors)
        if errors:
            raise ProjectValidationError(errors)
        connection.commit()
        for index, text in connection.execute("SELECT seq,payload FROM relations ORDER BY seq"):
            errors = []
            relation = normalize_graph_relation(json.loads(text), index, payload["schema"]["relation_types"], node_types, errors)
            if errors:
                raise ProjectValidationError(errors[:100])
            connection.execute(
                "UPDATE relations SET type=?, source_type=?, target_type=?, payload=? WHERE seq=?",
                (relation["type"], node_types[relation["source"]], node_types[relation["target"]], _dumps(relation), index),
            )
            if index % 1000 == 0:
                connection.commit()
                progress("校验关系", counts["nodes"], index + 1)
        connection.executescript("""
            CREATE INDEX nodes_group ON nodes(labels);
            CREATE INDEX relations_group ON relations(type, source_type, target_type);
        """)
        connection.commit()
        payload["counts"].update(counts)
        payload["counts"]["map_pois"] = spatial.count
        payload["graph_data"] = {"staging_path": str(staging)}
        payload["file_sha256"] = reader.sha256.hexdigest()
        progress("校验完成", counts["nodes"], counts["relations"])
        return payload
    except (ValueError, ijson.JSONError) as exc:
        if isinstance(exc, ProjectValidationError):
            raise
        raise ProjectValidationError([f"图谱 JSON 校验失败：{str(exc)[:1000]}"]) from exc
    finally:
        node_types.lookup.cache_clear()
        connection.close()


def _record_batches(cursor):
    batch, size = [], 0
    for (text,) in cursor:
        if batch and (len(batch) >= 250 or size + len(text) > 2 * 1024 * 1024):
            yield batch
            batch, size = [], 0
        batch.append(json.loads(text))
        size += len(text)
    if batch:
        yield batch


def import_staged_graph(graph_name: str, graph_data: Mapping,
                        progress: Progress = lambda *_: None) -> dict[str, int]:
    """Import only a private validated spool; never overwrite an existing graph."""
    db = _falkor_client()
    created = False
    done_nodes = done_relations = 0
    connection = None
    try:
        connection = sqlite3.connect(f"file:{Path(graph_data['staging_path']).as_posix()}?mode=ro", uri=True)
        connection.execute("PRAGMA cache_size=-4096")
        existing = {name.decode() if isinstance(name, bytes) else str(name) for name in db.list_graphs()}
        if graph_name in existing:
            raise FileExistsError(f"FalkorDB 图 {graph_name!r} 已存在")
        graph = db.select_graph(graph_name)
        for (labels_json,) in connection.execute("SELECT DISTINCT labels FROM nodes"):
            labels = json.loads(labels_json)
            expression = ":".join(_node_labels({"type": labels[0], "labels": labels}))
            cursor = connection.execute("SELECT payload FROM nodes WHERE labels=?", (labels_json,))
            for batch in _record_batches(cursor):
                rows = []
                for node in batch:
                    properties = _graph_safe_properties(node["properties"])
                    properties.setdefault("id", node["id"])
                    properties["__kg_node_id"] = node["id"]
                    rows.append({"properties": properties})
                created = True
                graph.query(f"UNWIND $rows AS row CREATE (n:{expression}) SET n += row.properties", {"rows": rows})
                done_nodes += len(batch)
                progress("导入节点", done_nodes, done_relations)
        for (label,) in connection.execute("SELECT DISTINCT type FROM nodes"):
            _node_labels({"type": label})
            graph.query(f"CREATE INDEX FOR (n:{label}) ON (n.__kg_node_id)")
        for signature in connection.execute("SELECT DISTINCT type,source_type,target_type FROM relations"):
            rel_type, source_type, target_type = signature
            for label in signature:
                _node_labels({"type": label})
            cursor = connection.execute(
                "SELECT payload FROM relations WHERE type=? AND source_type=? AND target_type=?", signature)
            for batch in _record_batches(cursor):
                rows = [{"source": row["source"], "target": row["target"],
                         "properties": _graph_safe_properties(row["properties"])} for row in batch]
                graph.query(
                    f"UNWIND $rows AS row MATCH (s:{source_type} {{__kg_node_id:row.source}}) "
                    f"MATCH (t:{target_type} {{__kg_node_id:row.target}}) "
                    f"CREATE (s)-[r:{rel_type}]->(t) SET r += row.properties", {"rows": rows})
                done_relations += len(batch)
                progress("导入关系", done_nodes, done_relations)
        return {
            "nodes": int(graph.query("MATCH (n) RETURN count(n)").result_set[0][0]),
            "relations": int(graph.query("MATCH ()-[r]->() RETURN count(r)").result_set[0][0]),
        }
    except FileExistsError:
        raise
    except Exception as exc:
        cleanup_error = None
        if created:
            try:
                db.select_graph(graph_name).delete()
            except Exception as cleanup_exc:
                cleanup_error = cleanup_exc
        raise GraphImportError(str(exc), cleanup_error=cleanup_error) from exc
    finally:
        if connection is not None:
            connection.close()
        db.close()