"""
cortex/brain_graph/builder.py — Brain Graph build pipeline.

Implements:
- BrainGraphBuilder: builds and refreshes brain graph indexes from
  file, directory, or URL sources.
- CompactNotationGenerator: generates ≤50-token compact representations
  for brain graph nodes.

Build pipeline (per BrainGraphBuilder.build()):
  1. Read source content
  2. Compute SHA-256 hash → dedup check
  3. Semantic chunking (respects code/markdown/text boundaries)
  4. Node extraction via LLM
  5. Edge extraction via LLM
  6. Compact notation generation per node
  7. Embed all nodes
  8. Store nodes, edges, embeddings
  9. Mark index as 'active'

All public methods are async.
"""
from __future__ import annotations

import hashlib
import json
import re
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from cortex.brain_graph.crud import BrainGraphCRUD
from cortex.config import CortexConfig
from cortex.llm.provider import LLMProvider
from cortex.storage.adapter import StorageAdapter
from cortex.types import BrainGraphEdge, BrainGraphIndex, BrainGraphNode

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_BUILDER_PROMPT_VER = "1.0"

# Node extraction JSON schema
_NODE_SCHEMA: Dict[str, Any] = {
    "type": "object",
    "properties": {
        "nodes": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "type": {"type": "string"},
                    "label": {"type": "string"},
                    "content": {"type": "string"},
                    "metadata": {"type": "object"},
                },
                "required": ["type", "label", "content"],
            },
        }
    },
    "required": ["nodes"],
}

# Edge extraction JSON schema
_EDGE_SCHEMA: Dict[str, Any] = {
    "type": "object",
    "properties": {
        "edges": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "source_label": {"type": "string"},
                    "target_label": {"type": "string"},
                    "relationship": {"type": "string"},
                    "weight": {"type": "number"},
                },
                "required": ["source_label", "target_label", "relationship"],
            },
        }
    },
    "required": ["edges"],
}

# Compact notation JSON schema
_COMPACT_SCHEMA: Dict[str, Any] = {
    "type": "object",
    "properties": {
        "compact": {"type": "string"}
    },
    "required": ["compact"],
}

# Chunk size limits
_MAX_CHUNK_CHARS = 1500
_MIN_CHUNK_CHARS = 100


# ---------------------------------------------------------------------------
# CompactNotationGenerator
# ---------------------------------------------------------------------------


class CompactNotationGenerator:
    """
    Generates token-efficient compact representations for brain graph nodes.

    Output is ≤50 tokens, injected into agent context windows during
    graph retrieval. Format is terse and information-dense.

    Example output for an API endpoint node:
        "POST /users → creates user. Requires: {name:str, email:str}.
         Returns: 201 {id, name, email}. Auth: Bearer. Rate: 100/min."
    """

    _SYSTEM = (
        "You generate ultra-compact, token-efficient summaries for knowledge graph nodes. "
        "Output must be ≤50 tokens. Be dense, specific, and drop filler words. "
        "Use terse notation: arrows (→), braces {}, colons (:). "
        "Capture the single most important fact about this node."
    )

    def __init__(self, llm: LLMProvider, model: Optional[str] = None) -> None:
        self._llm = llm
        self._model = model or llm.DEFAULT_COMPLETION_MODEL

    async def generate(
        self,
        node: BrainGraphNode,
        neighbors: Optional[List[BrainGraphNode]] = None,
    ) -> str:
        """
        Generate a compact representation for a node.

        Args:
            node: The node to summarize.
            neighbors: Optional adjacent nodes for context.

        Returns:
            A compact string ≤50 tokens.
        """
        neighbor_labels = ""
        if neighbors:
            labels = [n.label for n in neighbors[:5]]
            neighbor_labels = f"\nRelated: {', '.join(labels)}"

        user = (
            f"Node type: {node.node_type}\n"
            f"Label: {node.label}\n"
            f"Content: {node.content or '(none)'}{neighbor_labels}\n\n"
            "Generate a compact representation (≤50 tokens)."
        )

        from cortex.llm.router import resolve_routed_model
        _routed_model = resolve_routed_model(
            self._llm, "brain_graph_compact", self._model,
        )
        resp = await self._llm.extract(
            system=self._SYSTEM,
            user=user,
            schema=_COMPACT_SCHEMA,
            model=_routed_model,
            max_tokens=100,
        )
        compact = resp.data.get("compact", "")
        # Fallback: if LLM returns empty, use label as compact repr
        return compact.strip() if compact.strip() else node.label


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _compute_hash(content: str) -> str:
    """Return SHA-256 hex digest of content string."""
    return hashlib.sha256(content.encode("utf-8")).hexdigest()


def _read_file(path: str) -> str:
    """Read text content from a file path."""
    return Path(path).read_text(encoding="utf-8", errors="replace")


def _read_directory(path: str) -> str:
    """
    Read all text files in a directory recursively, concatenated.

    Skips binary files and common non-text extensions.
    """
    skip_suffixes = {
        ".pyc", ".pyo", ".so", ".dll", ".exe", ".bin", ".png",
        ".jpg", ".jpeg", ".gif", ".svg", ".ico", ".pdf", ".zip",
        ".tar", ".gz", ".whl", ".egg",
    }
    parts: List[str] = []
    for p in sorted(Path(path).rglob("*")):
        if p.is_file() and p.suffix.lower() not in skip_suffixes:
            try:
                parts.append(f"# File: {p}\n\n{p.read_text(encoding='utf-8', errors='replace')}")
            except OSError:
                pass
    return "\n\n---\n\n".join(parts)


def _read_url(url: str) -> str:
    """Fetch plain text content from a URL."""
    with urllib.request.urlopen(url, timeout=30) as resp:  # noqa: S310
        return resp.read().decode("utf-8", errors="replace")


def _read_source(source: str, source_type: str) -> str:
    """Dispatch source reading based on source_type."""
    if source_type == "file":
        return _read_file(source)
    elif source_type == "directory":
        return _read_directory(source)
    elif source_type == "url":
        return _read_url(source)
    else:
        raise ValueError(f"Unknown source_type: {source_type!r}. Expected 'file', 'directory', or 'url'.")


# ---------------------------------------------------------------------------
# Semantic chunker
# ---------------------------------------------------------------------------


def _semantic_chunk(content: str, max_chars: int = _MAX_CHUNK_CHARS) -> List[str]:
    """
    Split content into semantically meaningful chunks.

    Respects logical boundaries in this priority order:
    1. Markdown headers (## Header)
    2. Code fences (``` blocks)
    3. Double newlines (paragraph breaks)
    4. Single newlines (line breaks)
    5. Hard char limit as fallback

    Returns a list of non-empty chunk strings.
    """
    # Detect content type
    is_markdown = bool(re.search(r"^#{1,6}\s", content, re.MULTILINE))
    is_code = content.strip().startswith(("def ", "class ", "import ", "from ")) or \
              content.count("\n") > 5 and content.count("    ") > 2

    chunks: List[str] = []

    if is_markdown:
        # Split on markdown headers
        sections = re.split(r"(?=^#{1,6}\s)", content, flags=re.MULTILINE)
        for section in sections:
            section = section.strip()
            if not section:
                continue
            if len(section) <= max_chars:
                chunks.append(section)
            else:
                # Sub-split on double newlines
                chunks.extend(_split_by_paragraphs(section, max_chars))
    elif is_code:
        # Split code on function/class definitions
        chunks.extend(_split_code(content, max_chars))
    else:
        # Plain text: split on paragraphs
        chunks.extend(_split_by_paragraphs(content, max_chars))

    return [c for c in chunks if len(c.strip()) >= _MIN_CHUNK_CHARS]


def _split_by_paragraphs(text: str, max_chars: int) -> List[str]:
    """Split on double newlines, then hard-limit oversized paragraphs."""
    paragraphs = re.split(r"\n\n+", text)
    result: List[str] = []
    buffer = ""
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        if len(buffer) + len(para) + 2 <= max_chars:
            buffer = f"{buffer}\n\n{para}".strip() if buffer else para
        else:
            if buffer:
                result.append(buffer)
            if len(para) > max_chars:
                # Hard split on sentences
                result.extend(_hard_split(para, max_chars))
                buffer = ""
            else:
                buffer = para
    if buffer:
        result.append(buffer)
    return result


def _split_code(content: str, max_chars: int) -> List[str]:
    """Split code on top-level def/class boundaries."""
    lines = content.split("\n")
    chunks: List[str] = []
    current_lines: List[str] = []

    for line in lines:
        is_boundary = re.match(r"^(def |class |async def )", line)
        if is_boundary and current_lines:
            chunk = "\n".join(current_lines).strip()
            if chunk:
                if len(chunk) <= max_chars:
                    chunks.append(chunk)
                else:
                    chunks.extend(_hard_split(chunk, max_chars))
            current_lines = [line]
        else:
            current_lines.append(line)

    if current_lines:
        chunk = "\n".join(current_lines).strip()
        if chunk:
            if len(chunk) <= max_chars:
                chunks.append(chunk)
            else:
                chunks.extend(_hard_split(chunk, max_chars))

    return chunks


def _hard_split(text: str, max_chars: int) -> List[str]:
    """Last-resort split by character count at sentence boundaries."""
    chunks: List[str] = []
    while len(text) > max_chars:
        # Try to find a sentence boundary
        split_at = text.rfind(". ", 0, max_chars)
        if split_at == -1:
            split_at = max_chars
        else:
            split_at += 1  # include the period
        chunks.append(text[:split_at].strip())
        text = text[split_at:].strip()
    if text:
        chunks.append(text)
    return chunks


# ---------------------------------------------------------------------------
# BrainGraphBuilder
# ---------------------------------------------------------------------------


class BrainGraphBuilder:
    """
    Build and refresh brain graph indexes from source material.

    Usage:
        builder = BrainGraphBuilder(storage, config, llm)
        index = await builder.build(
            source="path/to/file.md",
            source_type="file",
            name="My Docs",
            owner_id="default",
        )

    The build pipeline:
    1. Read source content
    2. Compute SHA-256 hash → dedup (skip if hash matches existing active graph)
    3. Semantic chunk content
    4. Extract nodes per chunk via LLM
    5. Extract edges between nodes via LLM
    6. Generate compact notation per node
    7. Embed all nodes
    8. Store nodes, edges, embeddings
    9. Mark graph as active
    """

    _NODE_SYSTEM = (
        "You are a knowledge graph builder. Extract structured nodes from the given text. "
        "Node types: concept, endpoint, parameter, rule, example, definition, entity, procedure. "
        "Each node must have: type, label (short identifier), content (key information). "
        "Be precise. Extract only meaningful, distinct concepts — not filler. "
        "Return valid JSON matching the schema."
    )

    _EDGE_SYSTEM = (
        "You are a knowledge graph builder. Given a list of nodes extracted from a document, "
        "identify relationships between them. "
        "Relationship types: RELATED_TO, DEPENDS_ON, IMPLEMENTS, EXTENDS, CALLS, "
        "REFERENCES, EXAMPLES, CONTRASTS_WITH, PART_OF, SUPERSEDES. "
        "Only include edges where you are confident a real relationship exists. "
        "Weight 0.5-1.0 (1.0 = strong/direct, 0.5 = weak/indirect). "
        "Return valid JSON matching the schema."
    )

    def __init__(
        self,
        storage: StorageAdapter,
        config: CortexConfig,
        llm: LLMProvider,
    ) -> None:
        self._storage = storage
        self._config = config
        self._llm = llm
        self._crud = BrainGraphCRUD(storage)
        self._compact_gen = CompactNotationGenerator(llm)

    async def build(
        self,
        source: str,
        source_type: str,
        name: str,
        owner_id: str = "default",
        config: Optional[Dict[str, Any]] = None,
    ) -> BrainGraphIndex:
        """
        Build a brain graph index from source material.

        Args:
            source: File path, directory path, or URL.
            source_type: 'file', 'directory', or 'url'.
            name: Human-readable name for this graph.
            owner_id: Owner identifier (default: 'default').
            config: Optional build config overrides.

        Returns:
            BrainGraphIndex with status='active'.

        Raises:
            ValueError: If source_type is invalid.
            OSError: If source file/directory cannot be read.
        """
        # Step 1: Read source
        content = _read_source(source, source_type)

        # Step 2: Compute hash → dedup check
        source_hash = _compute_hash(content)
        existing = await self._crud.find_by_hash(source_hash, owner_id)
        if existing is not None:
            return existing

        # Step 3: Create index record (status=building)
        graph = await self._crud.create_graph(
            name=name,
            source_type=source_type,
            owner_id=owner_id,
            source_uri=source,
            source_hash=source_hash,
            embedding_model_id=self._config.embedding_model,
            builder_prompt_ver=_BUILDER_PROMPT_VER,
            index_config=config,
        )

        try:
            # Step 4+5: Chunk → extract nodes and edges
            chunks = _semantic_chunk(content)
            all_nodes, all_edges = await self._extract_graph(graph.id, chunks)

            # Step 6: Generate compact notation per node
            await self._generate_compact_reprs(graph.id, all_nodes)

            # Step 7+8: Embed nodes and store embeddings
            await self._embed_nodes(graph.id, all_nodes)

            # Step 9: Mark graph active
            now = datetime.now(timezone.utc).isoformat()
            graph = await self._crud.activate_graph(
                graph_id=graph.id,
                last_indexed=now,
                source_hash=source_hash,
                embedding_model_id=self._config.embedding_model,
            )

        except Exception:
            # Mark as failed if build pipeline errors
            await self._crud.update_graph(graph.id, status="failed")
            raise

        return graph

    async def refresh(self, graph_id: str) -> BrainGraphIndex:
        """
        Re-build a graph if its source has changed.

        Compares source_hash of the current index against a fresh read.
        If hash matches → return existing (no-op).
        If hash changed → delete old graph, build from scratch.

        Args:
            graph_id: ID of the graph to refresh.

        Returns:
            The refreshed (or existing) BrainGraphIndex.

        Raises:
            ValueError: If graph_id is not found.
        """
        graph = await self._crud.get_graph(graph_id)
        if graph is None:
            raise ValueError(f"Brain graph not found: {graph_id}")

        if graph.source_uri is None:
            raise ValueError(
                f"Graph {graph_id} has no source_uri — cannot refresh."
            )

        content = _read_source(graph.source_uri, graph.source_type)
        new_hash = _compute_hash(content)

        if graph.source_hash == new_hash:
            # Source unchanged — return existing
            return graph

        # Source changed: delete old graph, rebuild
        name = graph.name
        owner_id = graph.owner_id
        source_uri = graph.source_uri
        source_type = graph.source_type

        await self._crud.delete_graph(graph_id)

        return await self.build(
            source=source_uri,
            source_type=source_type,
            name=name,
            owner_id=owner_id,
        )

    # ------------------------------------------------------------------ #
    # Internal pipeline stages
    # ------------------------------------------------------------------ #

    async def _extract_graph(
        self, graph_id: str, chunks: List[str]
    ) -> Tuple[List[BrainGraphNode], List[BrainGraphEdge]]:
        """
        Extract nodes from chunks, then edges between all nodes.

        Returns (nodes, edges) both stored in the database.
        """
        stored_nodes: List[BrainGraphNode] = []

        # Node extraction: process each chunk
        for chunk in chunks:
            extracted = await self._extract_nodes_from_chunk(chunk)
            for raw_node in extracted:
                node = await self._crud.upsert_node(
                    graph_id=graph_id,
                    node_type=raw_node.get("type", "concept"),
                    label=raw_node.get("label", ""),
                    content=raw_node.get("content"),
                    metadata=raw_node.get("metadata"),
                )
                stored_nodes.append(node)

        # Edge extraction: use all node labels for context
        stored_edges: List[BrainGraphEdge] = []
        if len(stored_nodes) >= 2:
            node_index = {n.label: n for n in stored_nodes}
            raw_edges = await self._extract_edges(stored_nodes)
            for raw_edge in raw_edges:
                src_label = raw_edge.get("source_label", "")
                tgt_label = raw_edge.get("target_label", "")
                src_node = node_index.get(src_label)
                tgt_node = node_index.get(tgt_label)
                if src_node and tgt_node and src_node.id != tgt_node.id:
                    edge = await self._crud.upsert_edge(
                        graph_id=graph_id,
                        source_id=src_node.id,
                        target_id=tgt_node.id,
                        edge_type=raw_edge.get("relationship", "RELATED_TO"),
                        weight=float(raw_edge.get("weight", 1.0)),
                    )
                    stored_edges.append(edge)

        return stored_nodes, stored_edges

    async def _extract_nodes_from_chunk(
        self, chunk: str
    ) -> List[Dict[str, Any]]:
        """Call LLM to extract nodes from a single chunk."""
        user = f"Extract knowledge graph nodes from the following text:\n\n{chunk}"
        from cortex.llm.router import resolve_routed_model
        _routed_model = resolve_routed_model(
            self._llm, "brain_graph_nodes", self._config.extraction_model,
        )
        resp = await self._llm.extract(
            system=self._NODE_SYSTEM,
            user=user,
            schema=_NODE_SCHEMA,
            model=_routed_model,
            max_tokens=2000,
        )
        return resp.data.get("nodes", [])

    async def _extract_edges(
        self, nodes: List[BrainGraphNode]
    ) -> List[Dict[str, Any]]:
        """Call LLM to extract edges between the given nodes."""
        node_list = "\n".join(
            f"- [{n.node_type}] {n.label}: {(n.content or '')[:200]}"
            for n in nodes[:50]  # cap to avoid exceeding context
        )
        user = f"Identify relationships between these knowledge graph nodes:\n\n{node_list}"
        from cortex.llm.router import resolve_routed_model
        _routed_model = resolve_routed_model(
            self._llm, "brain_graph_edges", self._config.extraction_model,
        )
        resp = await self._llm.extract(
            system=self._EDGE_SYSTEM,
            user=user,
            schema=_EDGE_SCHEMA,
            model=_routed_model,
            max_tokens=2000,
        )
        return resp.data.get("edges", [])

    async def _generate_compact_reprs(
        self, graph_id: str, nodes: List[BrainGraphNode]
    ) -> None:
        """Generate and store compact representations for all nodes."""
        for node in nodes:
            compact = await self._compact_gen.generate(node)
            await self._crud.update_node(node.id, compact_repr=compact)

    async def _embed_nodes(
        self, graph_id: str, nodes: List[BrainGraphNode]
    ) -> None:
        """Embed all node contents and store in embedding_index."""
        if not nodes:
            return

        # Use content for embedding; fall back to label if content empty
        texts = [n.content or n.label for n in nodes]

        # Batch embed (LLM provider handles batching internally)
        resp = await self._llm.embed(
            texts=texts,
            model=self._config.embedding_model,
            input_type="document",
        )

        # Store each embedding
        for node, embedding in zip(nodes, resp.embeddings):
            await self._storage.store_embedding(
                item_type="brain_graph_node",
                item_id=node.id,
                embedding=embedding,
                model_id=self._config.embedding_model,
            )
