"""
hcp.protocol — Full Holographic Context Protocol pipeline.

Implements the complete HCP workflow:
  1. Chunking:   Split raw text into ContextChunks
  2. Seed:       Compute the HolographicSeed (Phase 1–2)
  3. Broadcast:  Attach seed to every chunk (Phase 3)
  4. Retrieve:   Query-driven holographic attention + ranking

This is the Integration Prerequisite (AX5 / Hayat):
  Seed computation (İlim), classification (İrade), and broadcast (Kudret)
  are independent capabilities that require this protocol class to
  INTEGRATE them into a living, functional pipeline.

Framework Axioms:
  AX5:  Integration prerequisite — protocol integrates the three phases
  AX37: Holographic Structure — every chunk carries the whole's structure
  AX13: Ehadiyet — local reflection at every chunk
  AX12: Vahidiyet — global unity maintained via seed
  AX27: Transparent vessel — protocol mediates, doesn't originate meaning
  KV₆:  Seed omnipresence — every chunk must have a seed
  KV₇:  Independence — seed computation is independent of query processing

KV₇ compliance: This module imports ONLY from hcp.* and standard library.
"""

from __future__ import annotations

import re
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple, cast

from hcp.types import (
    AttentionProfile,
    ContentType,
    ContextChunk,
    HCPConfig,
    HolographicSeed,
    SeededChunk,
    SourceMetadata,
    StageResult,
    WorkflowStage,
    WorkflowState,
)
from hcp.seed import compute_seed, extract_keywords
from hcp.attention import flat_attention, holographic_attention


class HCPProtocol:
    """Full Holographic Context Protocol pipeline.

    Usage:
        protocol = HCPProtocol()
        protocol.ingest("Your long context text here...")
        results = protocol.query("What is the key fact?")

        # Or from pre-built chunks:
        protocol = HCPProtocol()
        protocol.ingest_chunks(chunks)
        results = protocol.query("question", keywords=("key", "fact"))

    The protocol maintains:
      - chunks:  The context split into structural units
      - seed:    The holographic seed (computed once, broadcast everywhere)
      - seeded:  Chunks with seeds attached

    Per KV₇, the seed is computed INDEPENDENTLY of any query.
    The seed captures GLOBAL structure; queries are processed separately.
    """

    def __init__(self, config: Optional[HCPConfig] = None) -> None:
        self.config = config or HCPConfig()
        self._chunks: List[ContextChunk] = []
        self._seed: Optional[HolographicSeed] = None
        self._seeded: List[SeededChunk] = []
        self._workflows: Dict[str, WorkflowState] = {}

    @property
    def chunks(self) -> List[ContextChunk]:
        """The context chunks."""
        return list(self._chunks)

    @property
    def seed(self) -> Optional[HolographicSeed]:
        """The holographic seed (None if no context ingested)."""
        return self._seed

    @property
    def seeded_chunks(self) -> List[SeededChunk]:
        """Chunks with holographic seeds attached."""
        return list(self._seeded)

    @property
    def is_ready(self) -> bool:
        """Whether the protocol has ingested context and is ready for queries."""
        return self._seed is not None and len(self._seeded) > 0

    # -------------------------------------------------------------------
    # Ingestion: text → chunks → seed → seeded chunks
    # -------------------------------------------------------------------

    def ingest(
        self,
        text: str,
        source_metadata: Optional[SourceMetadata] = None,
        *,
        append: bool = True,
    ) -> HolographicSeed:
        """Ingest raw text: chunk, classify, compute seed, broadcast.

        This is the full Phase 1–3 pipeline:
          Phase 1 (İlim):   Split text into chunks, classify each
          Phase 2 (İrade):  Compute the holographic seed
          Phase 3 (Kudret): Broadcast seed to every chunk

        Args:
            text: The raw context text to ingest.
            source_metadata: Optional metadata from the originating tool.
                When provided, every chunk created from *text* inherits
                this metadata, enabling hint-based classification in
                ``compute_seed`` (three-tier confidence bypass / boost).
            append: If True (default), new chunks are ADDED to existing
                context.  If False, all existing chunks are replaced.
                v2.2.0 changed the default from False to True to prevent
                silent data loss (the validate_stage auto-ingest bug).

        Returns:
            The computed HolographicSeed.
        """
        chunks = self._chunk_text(text, source_metadata=source_metadata)
        return self.ingest_chunks(chunks, append=append)

    def ingest_chunks(
        self,
        chunks: List[ContextChunk],
        *,
        append: bool = True,
    ) -> HolographicSeed:
        """Ingest pre-built chunks: classify, compute seed, broadcast.

        Args:
            chunks: List of ContextChunk objects.
            append: If True (default), new chunks are appended to existing
                context.  If False, all existing chunks are replaced.

        Returns:
            The computed HolographicSeed.
        """
        if append and self._chunks:
            # Fix position offsets so new chunks follow existing ones
            offset = max(c.position for c in self._chunks) + 1
            adjusted: List[ContextChunk] = []
            for c in chunks:
                adjusted.append(ContextChunk(
                    content=c.content,
                    position=c.position + offset,
                    content_type=c.content_type,
                    keywords=c.keywords,
                    source_metadata=c.source_metadata,
                ))
            self._chunks.extend(adjusted)
        else:
            self._chunks = list(chunks)

        # Phase 1–2: compute seed (classification happens inside compute_seed)
        self._seed = compute_seed(self._chunks, self.config)

        # Phase 3: broadcast seed to every chunk (AX37 + KV₆)
        self._seeded = self._broadcast_seed()

        return self._seed

    def _chunk_text(
        self,
        text: str,
        source_metadata: Optional[SourceMetadata] = None,
    ) -> List[ContextChunk]:
        """Split raw text into ContextChunks.

        Chunking strategy: split by paragraphs first, then by size limit.
        Each paragraph is a natural structural unit.

        Args:
            text: Raw text to chunk.
            source_metadata: If provided, every chunk inherits this metadata
                so downstream classification can use hint-based bypass/boost.
        """
        # Split by double newlines (paragraph boundaries)
        paragraphs = re.split(r"\n\s*\n", text.strip())

        chunks: List[ContextChunk] = []
        position = 0

        for para in paragraphs:
            para = para.strip()
            if not para:
                continue

            # If paragraph exceeds chunk_size, split further
            if len(para) > self.config.chunk_size:
                # Split by sentences
                sentences = re.split(r"(?<=[.!?])\s+", para)
                current = ""
                for sentence in sentences:
                    if len(current) + len(sentence) > self.config.chunk_size and current:
                        chunks.append(ContextChunk(
                            content=current.strip(),
                            position=position,
                            source_metadata=source_metadata,
                        ))
                        position += 1
                        current = sentence
                    else:
                        current = current + " " + sentence if current else sentence

                if current.strip():
                    chunks.append(ContextChunk(
                        content=current.strip(),
                        position=position,
                        source_metadata=source_metadata,
                    ))
                    position += 1
            else:
                chunks.append(ContextChunk(
                    content=para,
                    position=position,
                    source_metadata=source_metadata,
                ))
                position += 1

        return chunks

    def _broadcast_seed(self) -> List[SeededChunk]:
        """Broadcast the holographic seed to every chunk.

        Per AX37 (Holographic Architecture): every part carries the whole.
        Per KV₆ (Seed Omnipresence): no chunk may lack a seed.
        Per T14 (Ne Ayn Ne Gayr): the seed participates in global truth
        without being identical to it.

        Also computes:
          - local_importance: this chunk's importance per the seed
          - keyword_overlap: keywords shared with global seed
        """
        if self._seed is None:
            return []

        seeded: List[SeededChunk] = []
        global_kw_set = set(self._seed.global_keywords)

        for chunk in self._chunks:
            # Local importance from seed
            if chunk.position < len(self._seed.position_importance):
                local_imp = self._seed.position_importance[chunk.position]
            else:
                local_imp = 0.0

            # Keyword overlap between chunk and global seed
            chunk_kw_set = set(chunk.keywords)
            overlap = tuple(sorted(chunk_kw_set & global_kw_set))

            seeded.append(SeededChunk(
                chunk=chunk,
                seed=self._seed,
                local_importance=local_imp,
                keyword_overlap=overlap,
            ))

        return seeded

    # -------------------------------------------------------------------
    # Query: retrieve relevant chunks using holographic attention
    # -------------------------------------------------------------------

    def query(
        self,
        question: str,
        keywords: Optional[Tuple[str, ...]] = None,
        top_k: int = 5,
    ) -> List[Tuple[SeededChunk, float]]:
        """Query the context using holographic attention.

        Computes attention profiles for both flat and holographic modes,
        then returns the top-k chunks ranked by holographic attention.

        Args:
            question: The query text.
            keywords: Optional explicit keywords. If None, extracted from question.
            top_k: Number of top results to return.

        Returns:
            List of (SeededChunk, attention_weight) tuples, sorted by weight.
        """
        if not self.is_ready:
            raise RuntimeError("Protocol not ready — call ingest() first")

        if keywords is None:
            keywords = extract_keywords(question, top_k=10)

        n = len(self._seeded)

        # Extract per-position keywords for content-query matching
        per_pos_kw = [
            sc.chunk.keywords if sc.chunk.keywords
            else extract_keywords(sc.chunk.content, top_k=10)
            for sc in self._seeded
        ]

        # Holographic attention (seed-modulated)
        assert self._seed is not None  # guaranteed by is_ready check above
        holo = holographic_attention(
            n_positions=n,
            seed=self._seed,
            query_position=n - 1,  # Query comes after context
            query_keywords=keywords,
            per_position_keywords=per_pos_kw,
            config=self.config,
        )

        # Pair chunks with their attention weights
        scored: List[Tuple[SeededChunk, float]] = []
        for i, sc in enumerate(self._seeded):
            weight = holo.weights[i] if i < len(holo.weights) else 0.0
            scored.append((sc, weight))

        # Sort by weight descending
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

    def compare_attention(
        self,
        question: str,
        keywords: Optional[Tuple[str, ...]] = None,
    ) -> Tuple[AttentionProfile, AttentionProfile]:
        """Compare flat vs holographic attention for a query.

        Returns both profiles for visualization or analysis.

        Args:
            question: The query text.
            keywords: Optional explicit keywords.

        Returns:
            Tuple of (flat_profile, holographic_profile).
        """
        if not self.is_ready:
            raise RuntimeError("Protocol not ready — call ingest() first")

        if keywords is None:
            keywords = extract_keywords(question, top_k=10)

        n = len(self._seeded)

        # Extract per-position keywords for content-query matching
        per_pos_kw = [
            sc.chunk.keywords if sc.chunk.keywords
            else extract_keywords(sc.chunk.content, top_k=10)
            for sc in self._seeded
        ]

        flat = flat_attention(n, query_position=n - 1, config=self.config)
        assert self._seed is not None  # guaranteed by is_ready check above
        holo = holographic_attention(
            n_positions=n,
            seed=self._seed,
            query_position=n - 1,
            query_keywords=keywords,
            per_position_keywords=per_pos_kw,
            config=self.config,
        )

        return flat, holo

    # -------------------------------------------------------------------
    # E2: Multi-Step Workflow Chains (AX2–4 three-phase actualization)
    # -------------------------------------------------------------------

    def create_workflow(
        self,
        stages: List[WorkflowStage],
        workflow_id: Optional[str] = None,
    ) -> WorkflowState:
        """Create a new multi-step workflow.

        Each workflow follows the three-phase actualization pattern
        (AX2–4): stages define distinction (İlim), the LLM selects
        action (İrade), and advancing executes it (Kudret).

        Args:
            stages: Ordered list of WorkflowStage definitions.
            workflow_id: Optional ID; auto-generated if omitted.

        Returns:
            The newly created WorkflowState.

        Raises:
            ValueError: If stages list is empty.
        """
        if not stages:
            raise ValueError("Workflow requires at least one stage")
        wid = workflow_id or uuid.uuid4().hex[:12]
        now = datetime.now(timezone.utc).isoformat()
        state = WorkflowState(
            workflow_id=wid,
            stages=list(stages),
            current_index=0,
            results=[],
            created_at=now,
        )
        self._workflows[wid] = state
        return state

    def advance_workflow(
        self,
        workflow_id: str,
        output: str,
        success: bool = True,
        error: str = "",
    ) -> WorkflowState:
        """Advance a workflow by recording the current stage's result.

        The LLM calls this after executing a stage.  The result is
        stored and current_index is advanced.

        Args:
            workflow_id: ID of an existing workflow.
            output: The stage's output (LLM-generated text, etc.).
            success: Whether the stage completed without error.
            error: Error message if success is False.

        Returns:
            Updated WorkflowState.

        Raises:
            KeyError: If workflow_id not found.
            RuntimeError: If workflow is already complete.
        """
        state = self._workflows.get(workflow_id)
        if state is None:
            raise KeyError(f"Workflow '{workflow_id}' not found")
        if state.is_complete:
            raise RuntimeError(
                f"Workflow '{workflow_id}' is already complete"
            )

        stage = state.current_stage
        assert stage is not None  # guaranteed by not is_complete check above
        now = datetime.now(timezone.utc).isoformat()
        result = StageResult(
            stage_name=stage.name,
            output=output,
            success=success,
            error=error,
            timestamp=now,
        )
        state.results.append(result)
        state.current_index += 1
        return state

    def get_workflow(self, workflow_id: str) -> WorkflowState:
        """Retrieve a workflow by ID.

        Args:
            workflow_id: The workflow identifier.

        Returns:
            The WorkflowState.

        Raises:
            KeyError: If not found.
        """
        state = self._workflows.get(workflow_id)
        if state is None:
            raise KeyError(f"Workflow '{workflow_id}' not found")
        return state

    # -------------------------------------------------------------------
    # Diagnostics
    # -------------------------------------------------------------------

    def diagnostics(self) -> Dict[str, object]:
        """Return diagnostic information about the current state.

        Per AX57 (Transparency): every output must disclose its
        epistemic status.
        """
        if not self.is_ready:
            result: Dict[str, object] = {"status": "not_ready", "n_chunks": 0}
        else:
            seed = self._seed
            assert seed is not None  # guaranteed by is_ready being True
            result = {
                "status": "ready",
                "n_chunks": seed.n_chunks,
                "type_distribution": {
                    ct.value: round(seed.type_signature[i], 3)
                    for i, ct in enumerate(ContentType)
                },
                "n_instructions": len(seed.instruction_positions),
                "n_constraints": len(seed.constraint_positions),
                "n_entities": len(seed.entity_positions),
                "n_structural_links": len(seed.structural_links),
                "n_global_keywords": len(seed.global_keywords),
                "coverage_ratio": round(seed.coverage_ratio, 3),
                "seed_has_structure": seed.has_structure,
            }

        # E2: Workflow summary
        if self._workflows:
            wf_summary = {}
            for wid, ws in self._workflows.items():
                wf_summary[wid] = {
                    "n_stages": len(ws.stages),
                    "current_index": ws.current_index,
                    "is_complete": ws.is_complete,
                    "n_results": len(ws.results),
                }
            result["workflows"] = wf_summary

        return result

    # -------------------------------------------------------------------
    # Export / Import (E3 — Checkpoint)
    # -------------------------------------------------------------------

    _SCHEMA_VERSION = 2
    _SUPPORTED_SCHEMA_VERSIONS = {1, 2}

    def export_state(self) -> Dict[str, object]:
        """Export the complete HCP state as a JSON-serializable dict.

        The exported state includes:
          - schema_version (for forward compatibility)
          - config
          - chunks (raw ContextChunk list)
          - seed (HolographicSeed, if computed)
          - seeded_chunks (SeededChunk list, if computed)
          - workflows

        Per T14 (Ne Ayn Ne Gayr): the exported snapshot participates
        in the protocol's truth without being identical to it — it is
        a frozen projection, not the living state.
        """
        state: Dict[str, object] = {
            "schema_version": self._SCHEMA_VERSION,
            "config": self.config.to_dict(),
            "chunks": [c.to_dict() for c in self._chunks],
        }
        if self._seed is not None:
            state["seed"] = self._seed.to_dict()
        else:
            state["seed"] = None

        state["seeded_chunks"] = [sc.to_dict() for sc in self._seeded]

        if self._workflows:
            state["workflows"] = {
                wid: ws.to_dict() for wid, ws in self._workflows.items()
            }
        else:
            state["workflows"] = {}

        return state

    def import_state(self, state: Dict[str, object]) -> None:
        """Restore HCP state from a previously exported dict.

        Resets the protocol to the exact state captured by export_state().
        Schema version is checked for forward compatibility.

        Args:
            state: Dict produced by export_state().

        Raises:
            ValueError: If schema_version is unsupported.
        """
        version = state.get("schema_version", 0)
        if version not in self._SUPPORTED_SCHEMA_VERSIONS:
            raise ValueError(
                f"Unsupported schema version {version} "
                f"(supported: {sorted(self._SUPPORTED_SCHEMA_VERSIONS)})"
            )

        # Config
        self.config = HCPConfig.from_dict(
            cast(Dict[str, object], state.get("config", {})),
        )

        # Chunks
        self._chunks = [
            ContextChunk.from_dict(cd)
            for cd in cast(List[Dict[str, object]], state.get("chunks", []))
        ]

        # Seed
        seed_data = state.get("seed")
        if seed_data is not None:
            self._seed = HolographicSeed.from_dict(
                cast(Dict[str, Any], seed_data),
            )
        else:
            self._seed = None

        # Seeded chunks
        self._seeded = [
            SeededChunk.from_dict(sd)
            for sd in cast(List[Dict[str, object]], state.get("seeded_chunks", []))
        ]

        # Workflows
        self._workflows = {}
        for wid, wd in cast(
            Dict[str, Dict[str, object]], state.get("workflows", {}),
        ).items():
            self._workflows[wid] = WorkflowState.from_dict(wd)
