import os
from typing import Any, Dict, List

from sayou.assembler import AssemblerPipeline
from sayou.chunking import ChunkingPipeline
from sayou.connector import ConnectorPipeline
from sayou.core.decorators import measure_time
from sayou.core.schemas import SayouBlock, SayouNode, SayouOutput
from sayou.loader import LoaderPipeline
from sayou.wrapper import WrapperPipeline

from ..core.base_brain import BaseBrainPipeline


class StructurePipeline(BaseBrainPipeline):
    """
    Structural assembly pipeline for code knowledge graphs.

    Transforms raw code files into a fully linked knowledge graph and
    persists it to the target store.

    Flow
    ────
    Connector → Chunking → Wrapper → Assembler → Loader

    Use cases
    ─────────
    - Code knowledge graph construction
    - RAG index over a codebase
    - Dependency / call-graph visualisation
    """

    component_name = "StructurePipeline"

    def __init__(
        self,
        extra_generators=None,
        extra_fetchers=None,
        extra_splitters=None,
        extra_adapters=None,
        extra_builders=None,
        extra_writers=None,
        config: Dict[str, Any] = None,
        **kwargs,
    ):
        super().__init__()

        self.config = self._init_config(config, kwargs)
        cfg = self.config

        self.connector = ConnectorPipeline(
            extra_generators=extra_generators,
            extra_fetchers=extra_fetchers,
            **cfg.connector,
        )
        self.chunking = ChunkingPipeline(
            extra_splitters=extra_splitters,
            **cfg.chunking,
        )
        self.wrapper = WrapperPipeline(
            extra_adapters=extra_adapters,
            **cfg.wrapper,
        )
        self.assembler = AssemblerPipeline(
            extra_builders=extra_builders,
            **cfg.assembler,
        )
        self.loader = LoaderPipeline(
            extra_writers=extra_writers,
            **cfg.loader,
        )

        self._sub_pipelines = {
            "connector": self.connector,
            "chunking": self.chunking,
            "wrapper": self.wrapper,
            "assembler": self.assembler,
            "loader": self.loader,
        }

        self.initialize(**kwargs)

    @classmethod
    def process(
        cls,
        source: str,
        destination: str = None,
        strategies: Dict[str, str] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        One-line facade.

        Example::

            StructurePipeline.process(
                "./my_project/",
                destination="bolt://localhost:7687",
                strategies={"assembler": "CodeGraphBuilder"},
            )
        """
        return cls(**kwargs).ingest(source, destination, strategies, **kwargs)

    @measure_time
    def ingest(
        self,
        source: str,
        destination: str = None,
        strategies: Dict[str, str] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """
        Execute the structural assembly pipeline.

        Args:
            source: Source URI or directory path.
            destination: Target store URI / path.
            strategies: Per-stage strategy overrides
                        (keys: connector, chunking, wrapper, assembler, loader).
        """
        self._emit("on_start", input_data={"source": source, "type": "structure"})
        strategies = strategies or {}
        stats: Dict[str, Any] = {
            "extracted": 0,
            "chunks": 0,
            "nodes": 0,
            "edges": 0,
            "saved": 0,
        }
        run_config = {**self.config._config, **kwargs}

        # ── Phase 1: Extract ─────────────────────────────────────────
        self._log(f"[1/4] Extracting from '{source}'...")
        try:
            packets_gen = self.connector.run(
                source,
                strategy=strategies.get("connector", "auto"),
                **run_config,
            )
            if not packets_gen:
                self._log("No data returned by connector.", level="warning")
                return stats
        except Exception as exc:
            self._log(f"[1/4] Extraction failed: {exc}", level="error")
            self._emit("on_error", error=exc)
            return stats

        # ── Phase 2: Chunk + Wrap (streaming) ────────────────────────
        self._log("[2/4] Chunking and wrapping...")
        all_nodes: List[SayouNode] = []

        for i, packet in enumerate(packets_gen):
            if not packet.success:
                continue

            stats["extracted"] += 1
            input_doc = self._normalise_packet(packet, i)
            if input_doc is None:
                continue

            chunks = self.chunking.run(
                input_doc,
                strategy=strategies.get("chunking", "auto"),
                **run_config,
            )
            if not chunks:
                continue

            stats["chunks"] += len(chunks)

            wrapper_out = self.wrapper.run(
                chunks,
                strategy=strategies.get("wrapper", "auto"),
                **run_config,
            )
            if wrapper_out and wrapper_out.nodes:
                all_nodes.extend(wrapper_out.nodes)

        if not all_nodes:
            self._log("No nodes collected. Aborting assembly.", level="error")
            return stats

        stats["nodes"] = len(all_nodes)
        self._log(f"Collected {stats['nodes']} nodes.")

        # ── Phase 3: Assemble ─────────────────────────────────────────
        self._log("[3/4] Assembling knowledge graph...")
        try:
            assembly_result = self.assembler.run(
                SayouOutput(nodes=all_nodes),
                strategy=strategies.get("assembler", "auto"),
                **run_config,
            )
            if isinstance(assembly_result, dict):
                stats["nodes"] = len(assembly_result.get("nodes", []))
                stats["edges"] = len(assembly_result.get("edges", []))
                self._log(f"Assembled {stats['nodes']} nodes, {stats['edges']} edges.")
        except Exception as exc:
            self._log(f"[3/4] Assembly failed: {exc}", level="error")
            self._emit("on_error", error=exc)
            return stats

        # ── Phase 4: Load ─────────────────────────────────────────────
        self._log(f"[4/4] Saving to '{destination}'...")
        try:
            success = self.loader.run(
                input_data=assembly_result,
                destination=destination,
                strategy=strategies.get("loader", "auto"),
                **run_config,
            )
            if success:
                stats["saved"] = 1
                self._log("Structure pipeline complete.")
            else:
                self._log("Loader returned False.", level="error")
        except Exception as exc:
            self._log(f"[4/4] Load failed: {exc}", level="error")
            self._emit("on_error", error=exc)

        self._emit("on_finish", result_data=stats, success=True)
        return stats

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _normalise_packet(self, packet: Any, index: int):
        """
        Convert a raw Connector packet into a SayouBlock for the Chunking stage.

        Handles dict payloads (GitHub connector), SayouBlock pass-throughs,
        and raw bytes / strings.  Returns None if the packet cannot be used.
        """
        raw = packet.data

        if isinstance(raw, dict):
            content = raw.get("content", "")
            meta = raw.get("meta", {}).copy()

            # Ensure a usable source path
            bad_sources = (None, "", "github_code", "unknown")
            if meta.get("source") in bad_sources:
                meta["source"] = (
                    meta.get("file_path")
                    or meta.get("path")
                    or meta.get("title")
                    or f"fallback_file_{index}"
                )

            # Inject extension if absent
            if not meta.get("extension"):
                _, ext = os.path.splitext(meta.get("source", ""))
                if ext:
                    meta["extension"] = ext.lower()

            if isinstance(content, bytes):
                content = content.decode("utf-8", errors="ignore")

            return SayouBlock(content=content, metadata=meta, type="code")

        if hasattr(raw, "content"):
            if "source" not in raw.metadata:
                raw.metadata["source"] = f"unknown_block_{index}"
            return raw

        if isinstance(raw, (bytes, str)):
            content = (
                raw.decode("utf-8", errors="ignore") if isinstance(raw, bytes) else raw
            )
            meta: Dict[str, Any] = dict(packet.meta or {})
            if packet.task:
                meta["source"] = packet.task.uri
                meta.update(packet.task.meta or {})
            else:
                meta["source"] = f"local_file_{index}"

            if not meta.get("extension"):
                _, ext = os.path.splitext(meta["source"])
                if ext:
                    meta["extension"] = ext.lower()

            return SayouBlock(content=content, metadata=meta, type="code")

        self._log(f"Unhandled packet data type: {type(raw).__name__}", level="warning")
        return None
