#!/usr/bin/env python3
"""
Echo Memory Tools - Search, Retrieval, and Plan Critique.

Functions:
- search(query, k): Hybrid BM25 + vector search
- get_receipt(query): Find specific receipt by ID/path
- answer(query): Evidence + Result + Confidence format
- critique_plan(plan_text): Review plans against receipts and constraints

Index locations (default):
- index/chunks.parquet
- index/vectors.sqlite (or .faiss)
- index/lexical.sqlite
"""

import json
import hashlib
import sqlite3
import struct
import re
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Any, Optional, Tuple
import os

# Default paths
ROOT = Path(__file__).resolve().parents[3]  # helix-cdc root
INDEX_DIR = ROOT / "index"
RECEIPTS_DIR = ROOT / "receipts"


@dataclass
class Chunk:
    """A chunk from the index."""
    id: str
    text: str
    source_path: str
    source_type: str
    sha256: str = ""
    relevance_score: float = 0.0
    chunk_type: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class AnswerResult:
    """Structured answer with evidence."""
    evidence: List[Dict[str, str]]
    result: str
    confidence: str  # HIGH | MEDIUM | LOW
    query: str
    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class CritiqueResult:
    """Plan critique output."""
    query: str
    patch: Optional[str]  # Suggested edits
    missing_proofs: List[str]
    risks: List[Dict[str, str]]
    proven_vs_assumed: Dict[str, List[str]]
    rollback_plan: str
    smallest_test: str
    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


class MemoryIndex:
    """Manager for Echo memory indexes."""

    def __init__(
        self,
        index_dir: Optional[Path] = None,
        receipts_dir: Optional[Path] = None
    ):
        self.index_dir = Path(index_dir) if index_dir else INDEX_DIR
        self.receipts_dir = Path(receipts_dir) if receipts_dir else RECEIPTS_DIR

        self._chunks_df = None
        self._vector_index = None
        self._lexical_conn = None
        self._embed_model = None

    def _ensure_chunks(self):
        """Load chunks parquet if needed."""
        if self._chunks_df is not None:
            return

        chunks_path = self.index_dir / "chunks.parquet"
        if not chunks_path.exists():
            raise FileNotFoundError(f"Chunks index not found: {chunks_path}")

        try:
            import pyarrow.parquet as pq
            table = pq.read_table(chunks_path)
            self._chunks_df = table.to_pandas()
        except ImportError:
            raise ImportError("pyarrow required: pip install pyarrow")

    def _ensure_vectors(self):
        """Load vector index if needed."""
        if self._vector_index is not None:
            return

        # Try FAISS first, then SQLite
        faiss_path = self.index_dir / "vectors.faiss"
        sqlite_path = self.index_dir / "vectors.sqlite"

        if faiss_path.exists():
            try:
                import faiss
                self._vector_index = {
                    "type": "faiss",
                    "index": faiss.read_index(str(faiss_path)),
                    "chunk_ids": json.loads(faiss_path.with_suffix(".mapping.json").read_text())["chunk_ids"]
                }
                return
            except ImportError:
                pass

        if sqlite_path.exists():
            self._vector_index = {
                "type": "sqlite",
                "conn": sqlite3.connect(str(sqlite_path))
            }
            return

        raise FileNotFoundError(f"Vector index not found in {self.index_dir}")

    def _ensure_lexical(self):
        """Load lexical index if needed."""
        if self._lexical_conn is not None:
            return

        lexical_path = self.index_dir / "lexical.sqlite"
        if not lexical_path.exists():
            raise FileNotFoundError(f"Lexical index not found: {lexical_path}")

        self._lexical_conn = sqlite3.connect(str(lexical_path))

    def _ensure_embedder(self):
        """Load embedding model if needed."""
        if self._embed_model is not None:
            return

        try:
            from sentence_transformers import SentenceTransformer
            model_name = os.environ.get("ECHO_EMBED_MODEL", "all-MiniLM-L6-v2")
            self._embed_model = SentenceTransformer(model_name)
        except ImportError:
            raise ImportError("sentence-transformers required: pip install sentence-transformers")

    def _embed_query(self, query: str) -> List[float]:
        """Embed a query string."""
        self._ensure_embedder()
        embedding = self._embed_model.encode(query)
        return embedding.tolist()

    def _vector_search(self, query_embedding: List[float], k: int = 8) -> List[Tuple[str, float]]:
        """Search vector index."""
        self._ensure_vectors()

        if self._vector_index["type"] == "faiss":
            import numpy as np
            import faiss

            query_np = np.array([query_embedding]).astype('float32')
            faiss.normalize_L2(query_np)

            scores, indices = self._vector_index["index"].search(query_np, k)
            results = []
            for score, idx in zip(scores[0], indices[0]):
                if idx >= 0:
                    chunk_id = self._vector_index["chunk_ids"][idx]
                    results.append((chunk_id, float(score)))
            return results
        else:
            # SQLite brute force
            import numpy as np

            cur = self._vector_index["conn"].cursor()
            cur.execute("SELECT chunk_id, embedding, dim FROM vectors")

            query_np = np.array(query_embedding)
            query_norm = np.linalg.norm(query_np)

            results = []
            for chunk_id, blob, dim in cur.fetchall():
                embedding = np.array(struct.unpack(f'{dim}f', blob))
                score = np.dot(query_np, embedding) / (query_norm * np.linalg.norm(embedding) + 1e-9)
                results.append((chunk_id, float(score)))

            results.sort(key=lambda x: x[1], reverse=True)
            return results[:k]

    def _lexical_search(self, query: str, k: int = 8) -> List[Tuple[str, float]]:
        """Search lexical index."""
        self._ensure_lexical()
        cur = self._lexical_conn.cursor()

        results = []
        try:
            # Try FTS5 first
            cur.execute("""
                SELECT id, bm25(chunks_fts) as score
                FROM chunks_fts
                WHERE chunks_fts MATCH ?
                ORDER BY score
                LIMIT ?
            """, (query, k * 2))

            for row in cur.fetchall():
                results.append((row[0], -row[1]))  # Negate FTS5 scores
        except sqlite3.OperationalError:
            # Fall back to LIKE
            terms = re.findall(r'\b\w+\b', query.lower())
            if terms:
                like_clause = " OR ".join(["text LIKE ?" for _ in terms])
                like_params = [f"%{term}%" for term in terms]
                cur.execute(f"""
                    SELECT id FROM chunks
                    WHERE {like_clause}
                    LIMIT ?
                """, like_params + [k * 2])
                for row in cur.fetchall():
                    results.append((row[0], 0.5))

        return results[:k]

    def _get_chunk_by_id(self, chunk_id: str) -> Optional[Chunk]:
        """Get chunk details by ID."""
        self._ensure_chunks()

        row = self._chunks_df[self._chunks_df["id"] == chunk_id]
        if row.empty:
            return None

        row = row.iloc[0]
        return Chunk(
            id=row["id"],
            text=row["text"],
            source_path=row.get("source_path", ""),
            source_type=row.get("source_type", ""),
            sha256=row.get("sha256", ""),
            chunk_type=row.get("chunk_type", ""),
        )

    def search(self, query: str, k: int = 8) -> List[Chunk]:
        """
        Hybrid search: BM25 + vector rerank.

        Args:
            query: Search query
            k: Number of results

        Returns:
            List of Chunk objects with relevance scores
        """
        # Get candidates from both indexes
        query_embedding = self._embed_query(query)
        vector_results = self._vector_search(query_embedding, k * 2)
        lexical_results = self._lexical_search(query, k * 2)

        # Combine and rerank
        scores = {}
        for chunk_id, score in vector_results:
            scores[chunk_id] = scores.get(chunk_id, 0) + score * 0.7  # Weight vector higher

        for chunk_id, score in lexical_results:
            scores[chunk_id] = scores.get(chunk_id, 0) + score * 0.3  # Weight lexical lower

        # Sort by combined score
        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:k]

        # Fetch full chunks
        chunks = []
        for chunk_id, score in ranked:
            chunk = self._get_chunk_by_id(chunk_id)
            if chunk:
                chunk.relevance_score = score
                chunks.append(chunk)

        return chunks

    def get_receipt(self, query: str) -> Optional[Dict[str, Any]]:
        """
        Find a specific receipt by ID/path match.

        Args:
            query: Receipt ID, path fragment, or work order ID

        Returns:
            Receipt dict with path, or None if not found
        """
        # Try direct file match first
        query_clean = query.strip().replace("/", "_").replace("\\", "_")

        # Search patterns
        patterns = [
            f"**/{query}*.json",
            f"**/*{query}*.json",
            f"**/proof_packs/{query}*/**/*.json",
        ]

        for pattern in patterns:
            for p in self.receipts_dir.glob(pattern):
                if p.is_file():
                    try:
                        data = json.loads(p.read_text())
                        data["_path"] = str(p)
                        data["_sha256"] = hashlib.sha256(p.read_bytes()).hexdigest()
                        return data
                    except Exception:
                        continue

        # Search in index
        try:
            self._ensure_lexical()
            cur = self._lexical_conn.cursor()
            cur.execute("""
                SELECT source_path FROM chunks
                WHERE source_type = 'receipt' AND (
                    source_path LIKE ? OR text LIKE ?
                )
                LIMIT 1
            """, (f"%{query}%", f"%{query}%"))

            row = cur.fetchone()
            if row:
                receipt_path = Path(row[0])
                if receipt_path.exists():
                    data = json.loads(receipt_path.read_text())
                    data["_path"] = str(receipt_path)
                    data["_sha256"] = hashlib.sha256(receipt_path.read_bytes()).hexdigest()
                    return data
        except Exception:
            pass

        return None

    def answer(self, query: str, k: int = 8) -> AnswerResult:
        """
        Answer a query with evidence citations.

        Format:
        - Evidence: List of [source] "quote" citations
        - Result: Synthesized answer from evidence only
        - Confidence: HIGH | MEDIUM | LOW

        Args:
            query: User question
            k: Number of evidence chunks to retrieve

        Returns:
            AnswerResult with evidence, result, and confidence
        """
        chunks = self.search(query, k)

        # Build evidence list
        evidence = []
        for chunk in chunks:
            # Truncate text for display
            text = chunk.text[:300] + "..." if len(chunk.text) > 300 else chunk.text
            text = text.replace("\n", " ").strip()
            evidence.append({
                "source": chunk.source_path,
                "quote": text,
                "type": chunk.source_type,
                "score": round(chunk.relevance_score, 3),
            })

        # Determine confidence based on evidence quality
        if not chunks:
            confidence = "LOW"
            result = "No relevant evidence found."
        else:
            max_score = max(c.relevance_score for c in chunks)
            receipt_count = sum(1 for c in chunks if c.source_type == "receipt")

            if max_score > 0.8 and receipt_count >= 2:
                confidence = "HIGH"
            elif max_score > 0.5 or receipt_count >= 1:
                confidence = "MEDIUM"
            else:
                confidence = "LOW"

            # Synthesize result from evidence
            if receipt_count > 0:
                receipt_chunks = [c for c in chunks if c.source_type == "receipt"]
                result = f"Based on {receipt_count} receipt(s) and {len(chunks) - receipt_count} other source(s)."

                # Extract key claims from receipts
                claims = []
                for c in receipt_chunks[:3]:
                    if "STATUS:" in c.text:
                        status_match = re.search(r"STATUS:\s*(\w+)", c.text)
                        if status_match:
                            claims.append(f"Status: {status_match.group(1)}")
                    if "CONCLUSION:" in c.text:
                        concl_match = re.search(r"CONCLUSION:\s*(.+?)(?:\n|$)", c.text)
                        if concl_match:
                            claims.append(f"Conclusion: {concl_match.group(1)[:100]}")

                if claims:
                    result += " Key findings: " + "; ".join(claims[:3])
            else:
                result = f"Found {len(chunks)} relevant chunks but no receipts. Consider verifying with receipts."

        return AnswerResult(
            evidence=evidence,
            result=result,
            confidence=confidence,
            query=query,
        )

    def critique_plan(self, plan_text: str) -> CritiqueResult:
        """
        Critique a plan against receipts and constraints.

        Checklist:
        - What is PROVEN vs ASSUMED?
        - What receipts will be produced?
        - What's the smallest test?
        - What could silently break?
        - What's the rollback?

        Args:
            plan_text: The plan to critique

        Returns:
            CritiqueResult with analysis
        """
        # Extract key terms from plan
        plan_lower = plan_text.lower()

        # Check for risky patterns
        risks = []

        risk_patterns = [
            (r"delete|remove|drop|reset", "Destructive operation", "HIGH"),
            (r"force|override|bypass|skip", "Bypassing safeguards", "MEDIUM"),
            (r"all|every|entire|whole", "Broad scope - consider incremental", "LOW"),
            (r"assume|probably|should work", "Unverified assumption", "MEDIUM"),
            (r"later|eventually|todo", "Deferred work", "LOW"),
        ]

        for pattern, desc, severity in risk_patterns:
            if re.search(pattern, plan_lower):
                risks.append({"pattern": pattern, "description": desc, "severity": severity})

        # Search for related receipts
        related_chunks = self.search(plan_text[:500], k=5)

        # Categorize proven vs assumed
        proven = []
        assumed = []

        proven_keywords = ["PROVEN", "VERIFIED", "PASS", "status.*true", "sha256"]
        assumed_keywords = ["planned", "will", "should", "TODO", "assumes"]

        for chunk in related_chunks:
            text_lower = chunk.text.lower()
            if any(re.search(kw, text_lower) for kw in proven_keywords):
                proven.append(f"[{chunk.source_path}] {chunk.text[:100]}...")
            if any(kw in text_lower for kw in assumed_keywords):
                assumed.append(f"[{chunk.source_path}] {chunk.text[:100]}...")

        # Identify missing proofs
        missing_proofs = []
        if "receipt" not in plan_lower:
            missing_proofs.append("No receipt generation mentioned")
        if "test" not in plan_lower and "verify" not in plan_lower:
            missing_proofs.append("No verification step mentioned")
        if "sha256" not in plan_lower and "hash" not in plan_lower:
            missing_proofs.append("No cryptographic verification mentioned")

        # Generate suggestions
        patch = None
        if missing_proofs:
            patch = "Consider adding:\n"
            if "No receipt generation mentioned" in missing_proofs:
                patch += "- Add receipt generation with SHA256 verification\n"
            if "No verification step mentioned" in missing_proofs:
                patch += "- Add verification test before/after\n"
            if "No cryptographic verification mentioned" in missing_proofs:
                patch += "- Add SHA256 hash checks for artifacts\n"

        # Smallest test suggestion
        smallest_test = "Unknown - no test specified in plan"
        if "test" in plan_lower:
            test_match = re.search(r"test[^.]*\.", plan_text, re.IGNORECASE)
            if test_match:
                smallest_test = test_match.group(0)

        # Rollback plan
        rollback = "No rollback mentioned"
        if "rollback" in plan_lower or "revert" in plan_lower:
            rollback_match = re.search(r"(rollback|revert)[^.]*\.", plan_text, re.IGNORECASE)
            if rollback_match:
                rollback = rollback_match.group(0)

        return CritiqueResult(
            query=plan_text[:200] + "..." if len(plan_text) > 200 else plan_text,
            patch=patch,
            missing_proofs=missing_proofs,
            risks=risks,
            proven_vs_assumed={
                "proven": proven[:5],
                "assumed": assumed[:5],
            },
            rollback_plan=rollback,
            smallest_test=smallest_test,
        )


# Module-level convenience functions
_default_index: Optional[MemoryIndex] = None


def _get_index() -> MemoryIndex:
    """Get or create default index."""
    global _default_index
    if _default_index is None:
        _default_index = MemoryIndex()
    return _default_index


def search(query: str, k: int = 8) -> List[Chunk]:
    """Search memory index."""
    return _get_index().search(query, k)


def get_receipt(query: str) -> Optional[Dict[str, Any]]:
    """Find a receipt by ID/path."""
    return _get_index().get_receipt(query)


def answer(query: str, k: int = 8) -> AnswerResult:
    """Answer with evidence and confidence."""
    return _get_index().answer(query, k)


def critique_plan(plan_text: str) -> CritiqueResult:
    """Critique a plan against receipts."""
    return _get_index().critique_plan(plan_text)
