# src/semantic_doc_search/core.py
#
# Fixes applied:
#   1. Action format — normalised to accept BOTH:
#        • Legacy:  'Ask doc-search "query"'       (backwards compat)
#        • Clean:   action='vector.search', context={'query': '...'}
#   2. Ingest separated from search — documents are embedded ONCE,
#      stored in the vector backend, and reused across queries.
#   3. Glob pattern fixed — pathlib doesn't support {txt,md} brace expansion.
#   4. vector.search + vector.insert actions added to match olang-resolver.json.

import os
import hashlib
from pathlib import Path
from typing import List, Dict, Any
import numpy as np
from sentence_transformers import SentenceTransformer

# ── Optional pgvector ─────────────────────────────────────────────────────────

try:
    import psycopg2
    from pgvector.psycopg2 import register_vector
    PGVECTOR_AVAILABLE = True
except ImportError:
    PGVECTOR_AVAILABLE = False
    print("⚠️  pgvector not available — using in-memory storage")

# ── Supported file extensions ─────────────────────────────────────────────────

SUPPORTED_EXTENSIONS = {".txt", ".md", ".pdf", ".html", ".docx"}


class SemanticDocSearch:
    def __init__(self, context: Dict[str, Any]):
        self.context     = context
        self.model_name  = "all-MiniLM-L6-v2"
        self.model       = SentenceTransformer(self.model_name)
        self.dimensions  = 384

        self.use_pgvector    = bool(context.get("POSTGRES_URL"))
        self.pg_conn         = None
        self.in_memory_store: List[Dict[str, Any]] = []

        if self.use_pgvector:
            if not PGVECTOR_AVAILABLE:
                raise ImportError(
                    "pgvector required for PostgreSQL support. "
                    "Install with: pip install pgvector psycopg2-binary"
                )
            self._init_pgvector()
        else:
            print("🔄 Using in-memory vector store")

    # ── pgvector init ─────────────────────────────────────────────────────────

    def _init_pgvector(self):
        try:
            self.pg_conn = psycopg2.connect(self.context["POSTGRES_URL"])
            register_vector(self.pg_conn)
            with self.pg_conn.cursor() as cur:
                cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS doc_embeddings (
                        id          TEXT PRIMARY KEY,
                        content     TEXT NOT NULL,
                        embedding   VECTOR(384),
                        source      TEXT,
                        metadata    JSONB DEFAULT '{}',
                        created_at  TIMESTAMPTZ DEFAULT NOW()
                    )
                """)
                self.pg_conn.commit()
            print("🗄️  pgvector initialised successfully")
        except Exception as e:
            print(f"❌ Failed to initialise pgvector: {e}")
            raise

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _hash_content(self, text: str) -> str:
        return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]

    def _embed_text(self, text: str) -> np.ndarray:
        embedding = self.model.encode(text, normalize_embeddings=True)
        if len(embedding) != self.dimensions:
            raise ValueError(
                f"Expected {self.dimensions}-dim embedding, got {len(embedding)}"
            )
        if np.allclose(embedding, 0.0):
            raise ValueError("Zero vector detected — embedding may have failed")
        return embedding

    def _chunk_text(
        self, text: str, chunk_size: int = 500, overlap: int = 50
    ) -> List[str]:
        if not text or not isinstance(text, str):
            return []
        chunk_size = max(1, chunk_size)
        overlap    = max(0, min(overlap, chunk_size // 2))
        chunks, start = [], 0
        while start < len(text):
            chunks.append(text[start : start + chunk_size])
            start += chunk_size - overlap
        return chunks

    # ── FIX 3: Correct glob — pathlib doesn't support {txt,md} ───────────────

    def _load_documents(self) -> List[Dict[str, Any]]:
        """
        Load documents from doc_root.
        Supports: .txt  .md  .pdf  .html  .docx
        FIX: iterate extensions individually — pathlib glob does NOT
             support bash-style brace expansion like *.{txt,md}.
        """
        doc_root = self.context.get("doc_root", "./docs")
        root_path = Path(doc_root)

        if not root_path.exists():
            print(f"⚠️  doc_root '{doc_root}' does not exist")
            return []

        docs = []

        # Iterate each supported extension separately
        for ext in SUPPORTED_EXTENSIONS:
            for file_path in root_path.glob(f"*{ext}"):
                if not file_path.is_file():
                    continue
                try:
                    content = self._extract_text(file_path)
                    if content.strip():
                        docs.append({
                            "id":      file_path.name,
                            "content": content,
                            "source":  f"file:{file_path.name}",
                        })
                        print(f"📄 Loaded: {file_path.name} ({len(content)} chars)")
                    else:
                        print(f"⚠️  Skipping empty file: {file_path.name}")
                except Exception as e:
                    print(f"❌ Failed to read {file_path.name}: {e}")

        return docs

    def _extract_text(self, file_path: Path) -> str:
        """Extract plain text from supported file types."""
        ext = file_path.suffix.lower()

        if ext in (".txt", ".md"):
            return file_path.read_text(encoding="utf-8")

        if ext == ".pdf":
            try:
                import pypdf
                reader = pypdf.PdfReader(str(file_path))
                return "\n".join(
                    page.extract_text() or "" for page in reader.pages
                )
            except ImportError:
                print("⚠️  pypdf not installed — skipping PDF. pip install pypdf")
                return ""

        if ext == ".html":
            try:
                from bs4 import BeautifulSoup
                soup = BeautifulSoup(file_path.read_text(encoding="utf-8"), "html.parser")
                return soup.get_text(separator=" ")
            except ImportError:
                print("⚠️  beautifulsoup4 not installed — skipping HTML. pip install beautifulsoup4")
                return ""

        if ext == ".docx":
            try:
                import docx
                doc = docx.Document(str(file_path))
                return "\n".join(p.text for p in doc.paragraphs)
            except ImportError:
                print("⚠️  python-docx not installed — skipping DOCX. pip install python-docx")
                return ""

        return ""

    # ── Upsert ────────────────────────────────────────────────────────────────

    def _upsert(self, doc_id: str, content: str, source: str, embedding: np.ndarray):
        if self.use_pgvector:
            self._upsert_pgvector(doc_id, content, source, embedding)
        else:
            self._upsert_memory(doc_id, content, source, embedding)

    def _upsert_pgvector(
        self, doc_id: str, content: str, source: str, embedding: np.ndarray
    ):
        with self.pg_conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO doc_embeddings (id, content, embedding, source)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE
                SET embedding = EXCLUDED.embedding,
                    content   = EXCLUDED.content
                """,
                (doc_id, content, embedding.tolist(), source),
            )
            self.pg_conn.commit()

    def _upsert_memory(
        self, doc_id: str, content: str, source: str, embedding: np.ndarray
    ):
        # Deduplicate by id
        self.in_memory_store = [
            d for d in self.in_memory_store if d["id"] != doc_id
        ]
        self.in_memory_store.append({
            "id":        doc_id,
            "content":   content,
            "embedding": embedding,
            "source":    source,
        })

    # ── Search ────────────────────────────────────────────────────────────────

    def _search(self, query_embedding: np.ndarray, top_k: int = 5) -> List[Dict]:
        if self.use_pgvector:
            return self._search_pgvector(query_embedding, top_k)
        return self._search_memory(query_embedding, top_k)

    def _search_pgvector(
        self, query_embedding: np.ndarray, top_k: int
    ) -> List[Dict]:
        with self.pg_conn.cursor() as cur:
            cur.execute(
                """
                SELECT id, content, source, embedding <-> %s AS distance
                FROM doc_embeddings
                ORDER BY distance
                LIMIT %s
                """,
                (query_embedding.tolist(), top_k),
            )
            return [
                {
                    "id":      r[0],
                    "content": r[1],
                    "source":  r[2],
                    "score":   round(1.0 - float(r[3]), 4),
                }
                for r in cur.fetchall()
            ]

    def _search_memory(
        self, query_embedding: np.ndarray, top_k: int
    ) -> List[Dict]:
        if not self.in_memory_store:
            return []
        from sklearn.metrics.pairwise import cosine_similarity

        scores = cosine_similarity(
            [query_embedding],
            [d["embedding"] for d in self.in_memory_store],
        )[0]

        ranked = sorted(
            zip(scores, self.in_memory_store), key=lambda x: x[0], reverse=True
        )
        return [
            {
                "id":      doc["id"],
                "content": doc["content"],
                "source":  doc["source"],
                "score":   round(float(score), 4),
            }
            for score, doc in ranked[:top_k]
        ]

    def _format_results(self, matches: List[Dict], query: str) -> Dict[str, Any]:
        return {
            "query":   query,
            "text":    "\n\n".join(m["content"] for m in matches),
            "matches": matches,
            "count":   len(matches),
        }

    # ─────────────────────────────────────────────────────────────────────────
    #  FIX 4: vector.insert and vector.search action handlers
    #  FIX 1: normalised action parsing — supports both legacy and clean format
    # ─────────────────────────────────────────────────────────────────────────

    def handle_action(self, action: str) -> Dict[str, Any]:
        """
        Route actions to the correct handler.

        Supported actions:
          vector.search               — search indexed documents (clean format)
          vector.insert               — ingest documents into vector store
          Ask doc-search "query"      — legacy format (backwards compat)

        Context fields for vector.search:
          query      (str)  required
          top_k      (int)  optional, default 5
          doc_root   (str)  optional — if set, auto-ingests before searching

        Context fields for vector.insert:
          doc_root   (str)  required — directory of documents to ingest
          chunk_size (int)  optional, default 500
          overlap    (int)  optional, default 50
        """

        # ── Clean action routing ──────────────────────────────────────────────
        if action == "vector.search":
            return self._handle_vector_search()

        if action == "vector.insert":
            return self._handle_vector_insert()

        # ── FIX 1: Legacy format  'Ask doc-search "query"' ───────────────────
        if action.startswith("Ask doc-search"):
            query = self._parse_legacy_query(action)
            return self._run_search(query)

        raise ValueError(
            f"Unknown action: '{action}'. "
            f"Supported: vector.search | vector.insert | Ask doc-search \"query\""
        )

    # ── Action implementations ────────────────────────────────────────────────

    def _handle_vector_search(self) -> Dict[str, Any]:
        """
        vector.search — search the index.

        If doc_root is provided AND the store is empty, auto-ingests first.
        This means a fresh deployment just works without a separate ingest call.
        """
        query = self.context.get("query", "").strip()
        if not query:
            raise ValueError("vector.search requires 'query' in context")

        top_k = int(self.context.get("top_k", 5))

        # Auto-ingest if store is empty and doc_root is provided
        store_empty = (
            not self.in_memory_store if not self.use_pgvector
            else not self._pgvector_has_docs()
        )
        if store_empty and self.context.get("doc_root"):
            print("🔄 Store empty — auto-ingesting documents before search...")
            self._ingest_documents()

        return self._run_search(query, top_k)

    def _handle_vector_insert(self) -> Dict[str, Any]:
        """
        vector.insert — ingest documents from doc_root into the vector store.
        Separates ingestion from search so documents are embedded ONCE.
        """
        if not self.context.get("doc_root"):
            raise ValueError("vector.insert requires 'doc_root' in context")

        count = self._ingest_documents()
        return {
            "inserted": count,
            "doc_root": self.context["doc_root"],
            "backend":  "pgvector" if self.use_pgvector else "memory",
        }

    def _ingest_documents(self) -> int:
        """
        FIX 2: Ingest is now SEPARATE from search.
        Documents are embedded once and stored — not re-processed on every query.
        """
        chunk_size = int(self.context.get("chunk_size", 500))
        overlap    = int(self.context.get("overlap",     50))
        documents  = self._load_documents()
        count      = 0

        print(f"🔄 Ingesting {len(documents)} documents...")

        for doc in documents:
            chunks = self._chunk_text(doc["content"], chunk_size, overlap) \
                     or [doc["content"]]

            for i, chunk in enumerate(chunks):
                if not chunk.strip():
                    continue
                try:
                    embedding = self._embed_text(chunk)
                    doc_id    = f"{doc['id']}:chunk:{i}"
                    self._upsert(doc_id, chunk, doc["source"], embedding)
                    count += 1
                except Exception as e:
                    print(f"⚠️  Skipping chunk {doc['id']}:{i} — {e}")

        print(f"✅ Ingested {count} chunks total")
        return count

    def _run_search(self, query: str, top_k: int = 5) -> Dict[str, Any]:
        """Execute the semantic search and return formatted results."""
        if not query.strip():
            return self._format_results([], query)

        print(f"🔍 Searching: '{query}' (top_k={top_k})")
        query_embedding = self._embed_text(query)
        matches         = self._search(query_embedding, top_k)

        print(f"📊 Found {len(matches)} matches")
        if matches:
            print(f"🎯 Top score: {matches[0]['score']:.4f}")

        return self._format_results(matches, query)

    def _pgvector_has_docs(self) -> bool:
        """Check if pgvector table has any rows."""
        if not self.pg_conn:
            return False
        with self.pg_conn.cursor() as cur:
            cur.execute("SELECT COUNT(*) FROM doc_embeddings LIMIT 1")
            return cur.fetchone()[0] > 0

    # ── FIX 1: Legacy query parser ────────────────────────────────────────────

    @staticmethod
    def _parse_legacy_query(action: str) -> str:
        """
        Parse query from legacy 'Ask doc-search "query"' format.
        Handles both quoted and unquoted variants.
        """
        if '"' in action:
            parts = action.split('"')
            if len(parts) >= 2:
                return parts[1].strip()

        # Fallback: everything after 'Ask doc-search'
        parts = action.split("Ask doc-search", 1)
        return parts[1].strip() if len(parts) > 1 else ""