from functools import lru_cache
from pathlib import Path

from sqlalchemy import cast, delete, func, or_, select

from .config import get_settings
from .knowledge_pages import page_text, relative_image_path
from .models import KnowledgeChunk, KnowledgeDocument, KnowledgePage


INDEX_VERSION = 2


def split_pages(pages: list[dict], *, size=650, overlap=100):
    for page in pages:
        text = page.get("text", "").strip()
        blocks = [text] if text else []
        for table in page.get("tables", []):
            blocks.append("表格: " + str(table.get("rows", table)))
        for block in blocks:
            start = 0
            while start < len(block):
                yield {"page": page["page_number"], "section": page.get("section", ""), "text": block[start:start + size]}
                if start + size >= len(block):
                    break
                start += size - overlap


@lru_cache(maxsize=2)
def local_encoder(model_path: str):
    path = Path(model_path)
    if not path.is_dir() or not any(path.iterdir()):
        raise RuntimeError(f"离线嵌入模型不存在: {path}；请挂载本地模型包")
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer(str(path), local_files_only=True, trust_remote_code=False)


def embed(texts: list[str], *, model_path: str | None = None):
    if not texts:
        return []
    model = local_encoder(model_path or str(get_settings().embedding_model_path))
    vectors = model.encode(texts, normalize_embeddings=True)
    output = []
    for vector in vectors:
        values = vector.tolist()
        if len(values) > 1024:
            raise RuntimeError("嵌入维度超过pgvector列容量1024")
        output.append(values + [0.0] * (1024 - len(values)))
    return output


def ingest(db, *, course_version_id: int, project_id: int | None, submission_file_id: int | None,
           title: str, category: str, source_hash: str, pages: list[dict], embedder=embed):
    existing = db.scalar(select(KnowledgeDocument).where(KnowledgeDocument.course_version_id == course_version_id,
                    KnowledgeDocument.project_id == project_id, KnowledgeDocument.submission_file_id == submission_file_id,
                    KnowledgeDocument.title == title))
    if existing and existing.source_hash == source_hash and existing.index_version == INDEX_VERSION and existing.status == "ready":
        return existing
    chunks = list(split_pages(pages))
    vectors = embedder([chunk["text"] for chunk in chunks])
    if len(vectors) != len(chunks):
        raise RuntimeError("嵌入向量数量与知识片段数量不一致")
    if existing:
        db.execute(delete(KnowledgePage).where(KnowledgePage.document_id == existing.id))
        db.execute(delete(KnowledgeChunk).where(KnowledgeChunk.document_id == existing.id))
        doc = existing
        doc.document_type, doc.source_hash = category, source_hash
    else:
        doc = KnowledgeDocument(course_version_id=course_version_id, project_id=project_id, submission_file_id=submission_file_id,
                                title=title, document_type=category, source_hash=source_hash)
        db.add(doc)
    doc.index_version, doc.status = INDEX_VERSION, "indexing"
    db.flush()
    for source in pages:
        db.add(KnowledgePage(document_id=doc.id, page=source["page_number"], text=page_text(source),
                             image_path=relative_image_path(source.get("rendered_image_path"))))
    for chunk, vector in zip(chunks, vectors):
        db.add(KnowledgeChunk(document_id=doc.id, course_version_id=course_version_id, project_id=project_id,
                              document_type=category, page=chunk["page"], section=chunk["section"], text=chunk["text"], embedding=vector))
    doc.status = "ready"
    db.commit()
    return doc


def retrieve(db, *, course_version_id: int, project_id: int, query: str, category: str | None = None,
             limit=10, embedder=embed, include_public=True, public_only=False):
    """Hybrid retrieval with project isolation and an optional public course namespace."""
    scope = (KnowledgeChunk.project_id.is_(None) if public_only else
             or_(KnowledgeChunk.project_id == project_id, KnowledgeChunk.project_id.is_(None)) if include_public else
             KnowledgeChunk.project_id == project_id)
    filters = [KnowledgeChunk.course_version_id == course_version_id, scope]
    if category:
        filters.append(KnowledgeChunk.document_type == category)
    vector = embedder([query])[0]
    if db.bind.dialect.name == "postgresql":
        from pgvector.sqlalchemy import Vector
        semantic = (1 - cast(KnowledgeChunk.embedding, Vector(1024)).cosine_distance(vector)).label("semantic")
        lexical = func.ts_rank(func.to_tsvector("simple", KnowledgeChunk.text), func.plainto_tsquery("simple", query)).label("lexical")
        rows = db.execute(select(KnowledgeChunk, semantic, lexical).where(*filters).order_by(semantic.desc(), lexical.desc()).limit(limit * 3)).all()
        ranked = [(row[0], float(row.semantic or 0) * .7 + float(row.lexical or 0) * .3) for row in rows]
    else:
        chunks = db.scalars(select(KnowledgeChunk).where(*filters)).all()
        terms = [query[i:i+2] for i in range(max(0, len(query)-1)) if query[i:i+2].strip()]
        ranked = []
        for chunk in chunks:
            semantic = sum(a*b for a, b in zip(chunk.embedding or [], vector))
            lexical = sum(term in chunk.text for term in terms) / max(len(terms), 1)
            ranked.append((chunk, semantic * .7 + lexical * .3))
    ranked.sort(key=lambda pair: pair[1], reverse=True)
    return [{"chunk_id": chunk.id, "document_id": chunk.document_id, "project_id": chunk.project_id,
             "document_type": chunk.document_type, "page": chunk.page, "section": chunk.section,
             "text": chunk.text, "rank": float(rank)} for chunk, rank in ranked[:limit]]