"""Qdrant REST adapter preserving the application's original chunk IDs and metadata.

The small collection interface matches the store's existing snapshot/rollback code.
REST calls have deadlines; backend errors never fall back to a stale local index.
"""
from __future__ import annotations
import hashlib
import math
import threading
import uuid
import httpx


class QdrantError(RuntimeError):
    def __init__(self, message, *, status_code=None):
        super().__init__(message)
        self.status_code = status_code


class QdrantClient:
    def __init__(self, url, api_key, prefix='yawan_', *, http_client=None):
        if not prefix or not all(c.isalnum() or c in '_-' for c in prefix):
            raise ValueError('Qdrant 知识库命名空间无效')
        self.prefix = prefix
        self._creation_lock = threading.Lock()
        self.http = http_client or httpx.Client(base_url=url.rstrip('/'),
            headers={'api-key': api_key}, timeout=httpx.Timeout(25, connect=5),
            limits=httpx.Limits(max_connections=8, max_keepalive_connections=4), trust_env=False)

    def request(self, method, path, body=None, *, missing_ok=False):
        try:
            response = self.http.request(method, path, json=body)
        except httpx.TimeoutException as exc:
            raise QdrantError('Qdrant 知识索引请求超时，请检查向量服务') from exc
        except httpx.HTTPError as exc:
            raise QdrantError('Qdrant 知识索引连接失败，请检查向量服务') from exc
        if missing_ok and response.status_code == 404:
            return None
        if not response.is_success:
            # No body: it can contain stored user text or reflected credentials.
            raise QdrantError(f'Qdrant 知识索引操作失败（HTTP {response.status_code}）', status_code=response.status_code)
        try:
            payload = response.json()
            return payload.get('result', payload)
        except (ValueError, AttributeError) as exc:
            raise QdrantError('Qdrant 知识索引返回了无效数据') from exc

    def get_collection(self, name):
        collection = QdrantCollection(self, name)
        result = self.request('GET', collection.path, missing_ok=True)
        if result is None:
            raise QdrantError(f'知识索引集合不存在：{name}')
        collection.configure(result)
        return collection

    def get_or_create_collection(self, name):
        collection = QdrantCollection(self, name)
        with self._creation_lock:
            description = self.request('GET', collection.path, missing_ok=True)
            if description is None:
                try:
                    self.request('PUT', collection.path, {
                        'vectors': {'size': 512, 'distance': 'Cosine'},
                        'on_disk_payload': True,
                        'optimizers_config': {'default_segment_number': 2},
                        'hnsw_config': {'m': 16, 'ef_construct': 128}})
                except QdrantError as exc:
                    # Another process can win after our GET; only accept a
                    # confirmed, compatible collection, never hide server errors.
                    if exc.status_code not in {400, 409}:
                        raise
                    description = self.request('GET', collection.path, missing_ok=True)
                    if description is None:
                        raise exc
                else:
                    description = {'config': {'params': {'vectors': {'size': 512, 'distance': 'Cosine'}}}}
                collection.configure(description)
                self.request('PUT', collection.path + '/index?wait=true',
                             {'field_name': 'meta.file', 'field_schema': 'keyword'})
            else:
                collection.configure(description)
        return collection

    def delete_collection(self, name):
        return self.request('DELETE', QdrantCollection(self, name).path, missing_ok=True)

    def close(self):
        self.http.close()


class QdrantCollection:
    def __init__(self, client, name):
        if not name or not all(c.isalnum() or c in '_-' for c in name):
            raise ValueError('知识索引集合名称无效')
        self.client, self.name = client, name
        self.path = '/collections/' + client.prefix + name
        self.distance = 'Cosine'

    def configure(self, description):
        try:
            vectors = description['config']['params']['vectors']
        except (KeyError, TypeError):
            vectors = None
        if not isinstance(vectors, dict) or vectors.get('size') != 512 or vectors.get('distance') not in {'Cosine', 'Euclid'}:
            raise QdrantError('知识索引集合必须使用 512 维 Cosine 或 Euclid 向量')
        self.distance = vectors['distance']

    def point_id(self, original_id):
        return str(uuid.uuid5(uuid.NAMESPACE_URL, self.client.prefix + self.name + ':' + str(original_id)))

    @staticmethod
    def filter(where):
        if not where:
            return None
        if not isinstance(where, dict) or set(where) != {'file'}:
            raise ValueError('不支持的知识索引过滤条件')
        value = where['file']
        if isinstance(value, dict):
            if set(value) != {'$in'} or not isinstance(value['$in'], (list, tuple, set)) or any(not isinstance(v, str) for v in value['$in']):
                raise ValueError('知识索引文件过滤条件无效')
            match = {'any': list(value['$in'])}
        elif isinstance(value, str):
            match = {'value': value}
        else:
            raise ValueError('知识索引文件过滤条件无效')
        return {'must': [{'key': 'meta.file', 'match': match}]}

    def count(self, where=None):
        body = {'exact': True}
        filtering = self.filter(where)
        if filtering is not None:
            body['filter'] = filtering
        result = self.client.request('POST', self.path + '/points/count', body)
        return int(result['count'])

    def _pages(self, where, include, size):
        if isinstance(size, bool) or not isinstance(size, int) or not 1 <= size <= 4000:
            raise ValueError('索引分页大小无效')
        offset = None
        seen_offsets = set()
        while True:
            body = {'limit': size, 'with_vector': 'embeddings' in include,
                    'with_payload': {'include': ['chunk_id', 'meta'] + (['text'] if 'documents' in include else [])}}
            filtering = self.filter(where)
            if filtering is not None:
                body['filter'] = filtering
            if offset is not None:
                body['offset'] = offset
            result = self.client.request('POST', self.path + '/points/scroll', body)
            yield result.get('points', [])
            new_offset = result.get('next_page_offset')
            if new_offset is None:
                break
            if new_offset in seen_offsets:
                raise QdrantError('Qdrant 分页游标未前进')
            seen_offsets.add(new_offset)
            offset = new_offset

    @staticmethod
    def _format(points, include):
        result = {'ids': [], 'documents': [], 'metadatas': [], 'embeddings': []}
        for point in points:
            payload = point.get('payload') or {}
            if not isinstance(payload.get('chunk_id'), str) or not payload['chunk_id']:
                raise QdrantError('知识索引原始片段编号缺失')
            result['ids'].append(payload['chunk_id'])
            if 'documents' in include:
                if not isinstance(payload.get('text'), str):
                    raise QdrantError('知识索引原文缺失')
                result['documents'].append(payload['text'])
            if 'metadatas' in include:
                if not isinstance(payload.get('meta'), dict):
                    raise QdrantError('知识索引来源元数据缺失')
                result['metadatas'].append(payload['meta'])
            if 'embeddings' in include:
                vector = point.get('vector')
                if not isinstance(vector, list) or len(vector) != 512 or any(not isinstance(v, (int, float)) or not math.isfinite(v) for v in vector):
                    raise QdrantError('知识索引向量维度无效')
                result['embeddings'].append(vector)
        return result

    def iter_records(self, where=None, include=None, page_size=256):
        include = ['documents', 'metadatas'] if include is None else include
        seen = set()
        for page in self._pages(where, include, page_size):
            self._check_scope(page, where)
            result = self._format(page, include)
            for identifier in result['ids']:
                if identifier in seen:
                    raise QdrantError('知识索引分页返回重复片段，请重新读取')
                seen.add(identifier)
            yield result

    @staticmethod
    def _check_scope(points, where):
        if not where:
            return
        allowed = where['file']
        for point in points:
            name = (point.get('payload', {}).get('meta') or {}).get('file')
            if not (name in allowed['$in'] if isinstance(allowed, dict) else name == allowed):
                raise QdrantError('知识索引返回了文件范围外的片段')

    def get(self, *, ids=None, where=None, include=None, limit=None, offset=0):
        include = ['documents', 'metadatas'] if include is None else include
        if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0 or (limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 0)):
            raise ValueError('索引分页范围无效')
        if limit == 0 or ids == []:
            return self._format([], include)
        if ids is not None:
            self.filter(where)
            found = {}
            for start in range(0, len(ids), 128):
                batch = ids[start:start + 128]
                points = self.client.request('POST', self.path + '/points',
                    {'ids': [self.point_id(identifier) for identifier in batch],
                     'with_payload': True, 'with_vector': 'embeddings' in include})
                for point in points:
                    original_id = (point.get('payload') or {}).get('chunk_id')
                    if original_id not in batch:
                        raise QdrantError('知识索引返回了请求范围外的片段')
                    if where:
                        name = (point.get('payload', {}).get('meta') or {}).get('file')
                        allowed = where['file']
                        if not (name in allowed['$in'] if isinstance(allowed, dict) else name == allowed):
                            continue
                    found[original_id] = point
            selected = [found[identifier] for identifier in ids if identifier in found]
            return self._format(selected[offset:None if limit is None else offset + limit], include)
        retained = []
        visited = 0
        for page in self._pages(where, include, min(256, max(1, limit or 256))):
            self._check_scope(page, where)
            for point in page:
                if visited >= offset:
                    retained.append(point)
                visited += 1
                if limit is not None and len(retained) >= limit:
                    return self._format(retained, include)
        return self._format(retained, include)

    def add(self, *, ids, documents, embeddings, metadatas):
        if not (len(ids) == len(documents) == len(embeddings) == len(metadatas)):
            raise ValueError('索引编号、正文、向量和元数据数量不一致')
        # Validate the entire input before the first batch is written.
        if any(not isinstance(identifier, str) or not identifier for identifier in ids) or len(set(ids)) != len(ids):
            raise ValueError('索引片段编号为空或重复')
        for text, vector, meta in zip(documents, embeddings, metadatas):
            if not isinstance(text, str) or not text.strip() or not isinstance(meta, dict):
                raise ValueError('索引正文或元数据无效')
            if not isinstance(vector, (list, tuple)) or len(vector) != 512 or any(not isinstance(v, (int, float)) or not math.isfinite(v) for v in vector):
                raise ValueError('索引向量必须是 512 维有限数值')
        for start in range(0, len(ids), 128):
            points = [{'id': self.point_id(ids[i]), 'vector': embeddings[i],
                       'payload': {'chunk_id': ids[i], 'text': documents[i], 'meta': metadatas[i]}}
                      for i in range(start, min(len(ids), start + 128))]
            self.client.request('PUT', self.path + '/points?wait=true', {'points': points})

    def delete(self, *, where=None, ids=None):
        if ids is not None and where is not None:
            raise ValueError('删除知识索引不可同时指定两种范围')
        if ids is not None:
            if not ids:
                return None
            body = {'points': [self.point_id(identifier) for identifier in ids]}
        elif where:
            body = {'filter': self.filter(where)}
        else:
            raise ValueError('删除知识索引必须指定文件或片段编号')
        return self.client.request('POST', self.path + '/points/delete?wait=true', body)

    def query(self, *, query_embeddings, n_results=6, where=None):
        if isinstance(n_results, bool) or not isinstance(n_results, int) or n_results < 1:
            raise ValueError('知识索引检索数量无效')
        result = {'ids': [], 'documents': [], 'metadatas': [], 'distances': []}
        for vector in query_embeddings:
            body = {'query': vector, 'limit': n_results, 'with_payload': True,
                    'with_vector': False, 'params': {'hnsw_ef': max(64, min(256, n_results * 2))}}
            filtering = self.filter(where)
            if filtering is not None:
                body['filter'] = filtering
            points = self.client.request('POST', self.path + '/points/query', body).get('points', [])
            self._check_scope(points, where)
            formatted = self._format(points, ['documents', 'metadatas'])
            for key in ('ids', 'documents', 'metadatas'):
                result[key].append(formatted[key])
            # Chroma's existing default is squared L2. Our embeddings are normalized.
            result['distances'].append([max(0., 2. * (1. - float(point['score']))) if self.distance == 'Cosine'
                                        else float(point['score']) ** 2 for point in points])
        return result


_client = None
_client_signature = None
_lock = threading.Lock()


def get_client():
    from .backend import qdrant_settings
    global _client, _client_signature
    url, key, prefix = qdrant_settings()
    signature = (url, hashlib.sha256(key.encode()).hexdigest(), prefix)
    with _lock:
        if _client is None or _client_signature != signature:
            if _client is not None:
                _client.close()
            _client = QdrantClient(url, key, prefix)
            _client_signature = signature
        return _client