"""Embedding service — encode text and manage vector storage."""

import hashlib
import json
import logging
import re
import time

_logger = logging.getLogger(__name__)

_model_instance = None

CHUNK_CHAR_LIMIT = 2000
CHUNK_OVERLAP = 200
_PARA_SPLIT = re.compile(r"\n\s*\n+")
_SENT_SPLIT = re.compile(r"(?<=[.!?])\s+")


def _chunk_text(text, *, limit=CHUNK_CHAR_LIMIT, overlap=CHUNK_OVERLAP):
    """Split long text into semantically-aware chunks of at most ``limit`` chars.

    Preserves paragraph and sentence boundaries where possible. Each chunk
    after the first retains ``overlap`` chars of tail context for continuity.
    Short inputs return a single-chunk list unchanged.
    """
    text = (text or "").strip()
    if not text:
        return []
    if len(text) <= limit:
        return [text]

    chunks = []
    buf = ""
    for para in _PARA_SPLIT.split(text):
        para = para.strip()
        if not para:
            continue
        candidate = (buf + "\n\n" + para) if buf else para
        if len(candidate) <= limit:
            buf = candidate
            continue
        if buf:
            chunks.append(buf)
            buf = ""
        if len(para) <= limit:
            buf = para
            continue
        for sent in _SENT_SPLIT.split(para):
            sent = sent.strip()
            if not sent:
                continue
            cand = (buf + " " + sent) if buf else sent
            if len(cand) <= limit:
                buf = cand
                continue
            if buf:
                chunks.append(buf)
            if len(sent) <= limit:
                buf = sent
            else:
                for i in range(0, len(sent), limit - overlap):
                    chunks.append(sent[i:i + limit])
                buf = ""
    if buf:
        chunks.append(buf)

    if overlap > 0 and len(chunks) > 1:
        with_overlap = [chunks[0]]
        for i in range(1, len(chunks)):
            tail = chunks[i - 1][-overlap:]
            with_overlap.append((tail + " " + chunks[i]).strip())
        chunks = with_overlap
    return chunks


def _sha256(text):
    return hashlib.sha256((text or "").encode("utf-8")).hexdigest()


def _get_model():
    """Lazy-load the sentence-transformers model (cached across calls)."""
    global _model_instance
    if _model_instance is None:
        try:
            from sentence_transformers import SentenceTransformer
            _model_instance = SentenceTransformer('all-MiniLM-L6-v2')
            _logger.info("Loaded sentence-transformers model: all-MiniLM-L6-v2")
        except ImportError:
            _logger.error(
                "sentence-transformers not installed. "
                "Run: pip install sentence-transformers"
            )
            raise
    return _model_instance


class EmbeddingService:
    """Encode texts, upsert embeddings, and perform semantic search."""

    def __init__(self, env):
        self.env = env
        self.Embedding = env['encoach.embedding'].sudo()

    def encode(self, texts):
        """Batch-encode texts to vectors.

        Args:
            texts: list of strings

        Returns:
            list of float lists (each 384-dim)
        """
        model = _get_model()
        embeddings = model.encode(texts, normalize_embeddings=True, show_progress_bar=False)
        return [e.tolist() for e in embeddings]

    def upsert(self, content_type, content_id, text, metadata=None):
        """Encode and store (or update) one embedding per text chunk.

        Long ``text`` is split by :func:`_chunk_text` into chunks of at most
        ``CHUNK_CHAR_LIMIT`` chars. Each chunk is stored as a separate row
        keyed by ``(content_type, content_id, chunk_index)``. Dedicated
        metadata columns (course_id, subject_id, entity_id, taxonomy,
        content_hash) are populated from ``metadata`` when present.

        Returns:
            list of encoach.embedding records (one per chunk). Empty list
            if input text is blank.
        """
        if not text or not text.strip():
            return []

        metadata = dict(metadata or {})
        chunks = _chunk_text(text)
        if not chunks:
            return []

        vectors = self.encode(chunks)
        existing_all = self.Embedding.search([
            ('content_type', '=', content_type),
            ('content_id', '=', content_id),
        ])
        by_idx = {rec.chunk_index: rec for rec in existing_all}
        stale = [rec for rec in existing_all if rec.chunk_index >= len(chunks)]
        if stale:
            self.Embedding.browse([r.id for r in stale]).unlink()

        records = []
        total = len(chunks)
        for idx, (chunk, vec) in enumerate(zip(chunks, vectors)):
            chunk_hash = _sha256(chunk)
            base_vals = {
                'content_text': chunk[:10000],
                'metadata_json': json.dumps(metadata),
                'course_id': int(metadata.get('course_id') or 0) or False,
                'subject_id': int(metadata.get('subject_id') or 0) or False,
                'entity_id': int(metadata.get('entity_id') or 0) or False,
                'taxonomy': metadata.get('taxonomy') or False,
                'content_hash': chunk_hash,
                'chunk_index': idx,
                'chunk_total': total,
            }
            rec = by_idx.get(idx)
            if rec:
                if rec.content_hash != chunk_hash or rec.chunk_total != total:
                    rec.write(base_vals)
                    rec.set_embedding(vec)
                else:
                    rec.write({k: v for k, v in base_vals.items()
                               if k in ('metadata_json', 'chunk_total')})
            else:
                rec = self.Embedding.create({
                    'content_type': content_type,
                    'content_id': content_id,
                    **base_vals,
                })
                rec.set_embedding(vec)
            records.append(rec)
        return records

    def search(self, query, *, content_type=None, limit=10,
               course_id=None, subject_id=None, entity_id=None,
               taxonomy=None):
        """Semantic search — encode query and find similar content.

        Optional metadata filters narrow the candidate pool to a specific
        course / subject / entity / taxonomy before ranking.

        Returns:
            list of dicts with text, metadata, similarity score
        """
        if not query or not query.strip():
            return []

        t0 = time.time()
        vectors = self.encode([query])
        results = self.Embedding.similarity_search(
            vectors[0],
            content_type=content_type,
            limit=limit,
            course_id=course_id,
            subject_id=subject_id,
            entity_id=entity_id,
            taxonomy=taxonomy,
        )
        latency = int((time.time() - t0) * 1000)
        _logger.info("Vector search for '%s' returned %d results in %dms",
                      query[:80], len(results), latency)
        return results

    def bulk_index(self, content_type, records_data):
        """Batch-index multiple records.

        Args:
            content_type: embedding content type
            records_data: list of dicts with keys: id, text, metadata
        """
        if not records_data:
            return 0

        texts = [r['text'] for r in records_data if r.get('text')]
        if not texts:
            return 0

        vectors = self.encode(texts)

        indexed = 0
        text_idx = 0
        for r in records_data:
            if not r.get('text'):
                continue
            self.upsert(content_type, r['id'], r['text'], r.get('metadata'))
            text_idx += 1
            indexed += 1

        _logger.info("Bulk-indexed %d %s records", indexed, content_type)
        return indexed

    def delete(self, content_type, content_id):
        """Remove an embedding."""
        existing = self.Embedding.search([
            ('content_type', '=', content_type),
            ('content_id', '=', content_id),
        ])
        if existing:
            existing.unlink()
