"""
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,
    ) -> 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).

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

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

        Args:
            chunks: List of ContextChunk objects.

        Returns:
            The computed HolographicSeed.
        """
        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 = 1

    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 != self._SCHEMA_VERSION:
            raise ValueError(
                f"Unsupported schema version {version} "
                f"(expected {self._SCHEMA_VERSION})"
            )

        # 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)
