"""Phase 1 Knowledge Base ingestion pipeline (scaffold)."""

from __future__ import annotations

import csv
import json
from pathlib import Path
from typing import Iterable, List, Optional

from .types import VectorMetadata, GraphNode, GraphEdge
from .adapters import (
    VectorStoreAdapter,
    GraphStoreAdapter,
    WeaviateAdapter,
    ChromaAdapter,
    Neo4jAdapter,
    SnowflakeGraphAdapter,
)
from src.config import settings
from src.datashield import mask_row


class KBIngestionPipeline:
    """Simple ingestion pipeline for demo and testing."""

    def __init__(self, base_dir: Path | None = None):
        self.base_dir = Path(base_dir) if base_dir else Path(settings.data_dir) / "knowledge_base"
        self.base_dir.mkdir(parents=True, exist_ok=True)
        self._prefer_file_backed = base_dir is not None
        self.vector_store = self._init_vector_store()
        self.graph_store = self._init_graph_store()

    def _init_vector_store(self):
        if self._prefer_file_backed:
            return VectorStoreAdapter(self.base_dir)
        backend = settings.kb_vector_backend.lower()
        try:
            if backend == "weaviate":
                return WeaviateAdapter(settings.weaviate_url, settings.weaviate_api_key)
            if backend == "chroma":
                return ChromaAdapter(str(settings.chroma_path))
        except Exception:
            return VectorStoreAdapter(self.base_dir)
        return VectorStoreAdapter(self.base_dir)

    def _init_graph_store(self):
        if self._prefer_file_backed:
            return GraphStoreAdapter(self.base_dir)
        backend = settings.kb_graph_backend.lower()
        try:
            if backend == "neo4j":
                return Neo4jAdapter(settings.neo4j_uri, settings.neo4j_user, settings.neo4j_password)
            if backend == "snowflake":
                return SnowflakeGraphAdapter(settings.snowflake_graph_database, settings.snowflake_graph_schema)
        except Exception:
            return GraphStoreAdapter(self.base_dir)
        return GraphStoreAdapter(self.base_dir)

    def ingest_vector_metadata(self, items: Iterable[VectorMetadata]) -> int:
        return self.vector_store.add_metadata(items)

    def ingest_graph(self, nodes: Iterable[GraphNode], edges: Iterable[GraphEdge]) -> int:
        self.graph_store.upsert_nodes(nodes)
        self.graph_store.upsert_edges(edges)
        return len(list(edges))

    def ingest_from_lineage_json(self, path: Path) -> int:
        data = json.loads(Path(path).read_text(encoding="utf-8"))
        nodes = [GraphNode(**n) for n in data.get("nodes", [])]
        edges = [GraphEdge(source=e["from"], target=e["to"], type=e["type"]) for e in data.get("edges", [])]
        return self.ingest_graph(nodes, edges)

    def ingest_from_csv(
        self,
        path: Path,
        source_type: str,
        id_column: str,
        hierarchy_id: Optional[str] = None,
        mask_pii: bool = False,
        can_view_pii: bool = False,
    ) -> int:
        items: List[VectorMetadata] = []
        with Path(path).open(newline="", encoding="utf-8") as f:
            for row in csv.DictReader(f):
                src_id = row.get(id_column, "")
                if not src_id:
                    continue
                if mask_pii:
                    row = mask_row(row, can_view_pii=can_view_pii)
                items.append(VectorMetadata(
                    id=f"{source_type}:{src_id}",
                    source_type=source_type,
                    source_id=src_id,
                    hierarchy_id=hierarchy_id,
                    tags=[source_type],
                    extra=row,
                ))
        return self.ingest_vector_metadata(items)
