from __future__ import annotations

from dataclasses import dataclass
from io import BytesIO, StringIO
import csv
import json
from pathlib import Path
import re
from typing import Iterable
import zipfile

import numpy as np
from sentence_transformers import SentenceTransformer


SUPPORTED_MODELS: list[tuple[str, str]] = [
    ("BAAI/bge-small-en-v1.5", "Lightweight English retrieval"),
    ("sentence-transformers/all-MiniLM-L6-v2", "Very easy lightweight baseline"),
    ("intfloat/e5-small-v2", "Small E5 family model"),
    ("mixedbread-ai/mxbai-embed-large-v1", "Higher quality larger model"),
    ("nomic-ai/nomic-embed-text-v1.5", "Longer-context friendly"),
]

SUMMARY_FILES = {"pooling_summary.json", "embedding_summary.json", "querying_summary.json"}
TEXT_SUFFIXES = {".md", ".txt"}
JSONL_SUFFIX = ".jsonl"
NPZ_SUFFIX = ".npz"


@dataclass
class PooledRecord:
    pooled_id: str
    group_key: str
    member_count: int
    embedding_dim: int
    source_model_name: str
    pooling_method: str
    embedding: np.ndarray


@dataclass
class MarkdownChunk:
    group_key: str
    source_file: str
    text: str
    char_count: int
    word_count: int


@dataclass
class QueryResult:
    rank: int
    score: float
    pooled_id: str
    group_key: str
    source_markdown_file: str
    member_count: int
    embedding_dim: int
    source_model_name: str
    query_model_name: str
    pooling_method: str
    char_count: int
    word_count: int
    preview: str
    text: str


class QueryingError(RuntimeError):
    pass


def supported_models() -> list[tuple[str, str]]:
    return SUPPORTED_MODELS.copy()


def infer_query_prefix(model_name: str) -> str:
    name = model_name.lower()
    if "e5" in name:
        return "query:"
    return ""


def _safe_stem(filename: str) -> str:
    return re.sub(r"[^A-Za-z0-9._-]+", "_", Path(filename).stem).strip("_") or "querying_output"


def _normalize_rows(mat: np.ndarray) -> np.ndarray:
    mat = np.asarray(mat, dtype=np.float32)
    if mat.ndim == 1:
        mat = mat.reshape(1, -1)
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    norms = np.where(norms == 0, 1.0, norms)
    return mat / norms


def _load_summary_from_zip(zf: zipfile.ZipFile) -> dict:
    for name in zf.namelist():
        if Path(name).name in SUMMARY_FILES:
            try:
                return json.loads(zf.read(name).decode("utf-8"))
            except Exception:
                return {}
    return {}


def _load_pooled_records_from_npz(npz_bytes: bytes, summary: dict) -> list[PooledRecord]:
    data = np.load(BytesIO(npz_bytes), allow_pickle=True)
    required = {"ids", "group_keys", "member_counts", "embeddings"}
    if not required.issubset(set(data.files)):
        raise QueryingError("The pooled NPZ file does not contain the expected arrays.")

    ids = data["ids"]
    group_keys = data["group_keys"]
    member_counts = data["member_counts"]
    embeddings = data["embeddings"]

    source_model_name = str(summary.get("source_model_name") or summary.get("model_name") or "unknown")
    pooling_method = str(summary.get("pooling_method", "unknown"))

    records: list[PooledRecord] = []
    for i in range(len(ids)):
        records.append(
            PooledRecord(
                pooled_id=str(ids[i]),
                group_key=str(group_keys[i]),
                member_count=int(member_counts[i]),
                embedding_dim=int(embeddings.shape[1]),
                source_model_name=source_model_name,
                pooling_method=pooling_method,
                embedding=np.asarray(embeddings[i], dtype=np.float32),
            )
        )
    return records


def _load_pooled_records_from_jsonl(jsonl_text: str, summary: dict) -> list[PooledRecord]:
    default_source_model_name = str(summary.get("source_model_name") or summary.get("model_name") or "unknown")
    default_pooling_method = str(summary.get("pooling_method", "unknown"))
    records: list[PooledRecord] = []
    for line in jsonl_text.splitlines():
        line = line.strip()
        if not line:
            continue
        obj = json.loads(line)
        emb = np.asarray(obj.get("embedding", []), dtype=np.float32)
        records.append(
            PooledRecord(
                pooled_id=str(obj.get("pooled_id") or obj.get("id") or ""),
                group_key=str(obj.get("group_key") or obj.get("source_file") or ""),
                member_count=int(obj.get("member_count", 1)),
                embedding_dim=int(obj.get("embedding_dim", emb.shape[0] if emb.ndim == 1 else 0)),
                source_model_name=str(obj.get("source_model_name", default_source_model_name)),
                pooling_method=str(obj.get("pooling_method", default_pooling_method)),
                embedding=emb,
            )
        )
    return records


def load_pooled_zip_bytes(zip_bytes: bytes) -> tuple[list[PooledRecord], dict]:
    with zipfile.ZipFile(BytesIO(zip_bytes), "r") as zf:
        summary = _load_summary_from_zip(zf)

        npz_name = next((name for name in zf.namelist() if name.lower().endswith(NPZ_SUFFIX)), None)
        if npz_name:
            records = _load_pooled_records_from_npz(zf.read(npz_name), summary)
            if records:
                return records, summary

        jsonl_name = next((name for name in zf.namelist() if name.lower().endswith(JSONL_SUFFIX)), None)
        if jsonl_name:
            records = _load_pooled_records_from_jsonl(zf.read(jsonl_name).decode("utf-8"), summary)
            if records:
                return records, summary

    raise QueryingError("No pooled embedding payload (.npz or .jsonl) was found in the uploaded zip.")


def load_markdown_zip_bytes(zip_bytes: bytes) -> dict[str, MarkdownChunk]:
    chunks: dict[str, MarkdownChunk] = {}
    with zipfile.ZipFile(BytesIO(zip_bytes), "r") as zf:
        for name in sorted(zf.namelist()):
            if name.endswith("/"):
                continue
            suffix = Path(name).suffix.lower()
            if suffix not in TEXT_SUFFIXES:
                continue
            text = zf.read(name).decode("utf-8", errors="ignore").strip()
            if not text:
                continue
            group_key = Path(name).stem
            chunks[group_key] = MarkdownChunk(
                group_key=group_key,
                source_file=Path(name).name,
                text=text,
                char_count=len(text),
                word_count=len(text.split()),
            )
    if not chunks:
        raise QueryingError("No .md or .txt chunk files were found in the markdown chunk zip.")
    return chunks


def _apply_prefix(query: str, prefix: str) -> str:
    prefix = prefix.strip()
    if not prefix:
        return query
    if query.startswith(prefix + " "):
        return query
    return f"{prefix} {query}"


def query_pooled_records(
    pooled_records: list[PooledRecord],
    markdown_chunks: dict[str, MarkdownChunk],
    query_text: str,
    model_name: str,
    top_k: int = 5,
    query_prefix: str = "",
    normalize_query: bool = True,
) -> tuple[list[dict], list[QueryResult], dict]:
    if not pooled_records:
        raise QueryingError("No pooled records were provided.")
    query_text = query_text.strip()
    if not query_text:
        raise QueryingError("The query text is empty.")

    model = SentenceTransformer(model_name)
    model_input = _apply_prefix(query_text, query_prefix)
    query_vector = model.encode([model_input], normalize_embeddings=normalize_query, show_progress_bar=False)
    query_vector = np.asarray(query_vector[0], dtype=np.float32)

    pooled_matrix = np.stack([r.embedding for r in pooled_records], axis=0).astype(np.float32)
    pooled_matrix = _normalize_rows(pooled_matrix)
    q = _normalize_rows(query_vector)[0]

    scores = pooled_matrix @ q
    order = np.argsort(-scores)
    top_k = max(1, min(int(top_k), len(order)))

    results: list[QueryResult] = []
    manifest: list[dict] = []
    for rank, idx in enumerate(order[:top_k], start=1):
        rec = pooled_records[int(idx)]
        score = float(scores[int(idx)])
        md = markdown_chunks.get(rec.group_key)
        source_markdown_file = md.source_file if md else ""
        text = md.text if md else ""
        char_count = md.char_count if md else 0
        word_count = md.word_count if md else 0
        preview = text[:220].replace("\n", " ") if text else ""
        result = QueryResult(
            rank=rank,
            score=score,
            pooled_id=rec.pooled_id,
            group_key=rec.group_key,
            source_markdown_file=source_markdown_file,
            member_count=rec.member_count,
            embedding_dim=rec.embedding_dim,
            source_model_name=rec.source_model_name,
            query_model_name=model_name,
            pooling_method=rec.pooling_method,
            char_count=char_count,
            word_count=word_count,
            preview=preview,
            text=text,
        )
        results.append(result)
        manifest.append(
            {
                "rank": rank,
                "score": score,
                "pooled_id": rec.pooled_id,
                "group_key": rec.group_key,
                "source_markdown_file": source_markdown_file,
                "member_count": rec.member_count,
                "embedding_dim": rec.embedding_dim,
                "source_model_name": rec.source_model_name,
                "query_model_name": model_name,
                "pooling_method": rec.pooling_method,
                "char_count": char_count,
                "word_count": word_count,
                "preview": preview,
            }
        )

    summary = {
        "query_text": query_text,
        "query_input_text": model_input,
        "query_model_name": model_name,
        "query_prefix": query_prefix,
        "top_k": top_k,
        "candidate_count": len(pooled_records),
        "best_score": float(results[0].score) if results else None,
        "source_model_name": pooled_records[0].source_model_name if pooled_records else "unknown",
        "pooling_method": pooled_records[0].pooling_method if pooled_records else "unknown",
        "markdown_chunk_count": len(markdown_chunks),
        "normalize_query": bool(normalize_query),
        "query_vector_dim": int(query_vector.shape[0]),
    }
    return manifest, results, summary


def build_llm_prompt(query_text: str, results: list[QueryResult]) -> str:
    parts: list[str] = []
    for result in results:
        if not result.text:
            continue
        parts.append(
            f"[Rank {result.rank} | Score {result.score:.4f} | File {result.source_markdown_file or result.group_key}]\n{result.text}"
        )
    context = "\n\n".join(parts)
    return (
        "Use only the context below to answer the question.\n"
        "If the answer is not in the context, say you do not know.\n\n"
        f"Context:\n{context}\n\n"
        f"Question:\n{query_text}\n"
    )


def _results_csv_bytes(manifest: list[dict]) -> bytes:
    out = StringIO()
    fieldnames = [
        "rank",
        "score",
        "pooled_id",
        "group_key",
        "source_markdown_file",
        "member_count",
        "embedding_dim",
        "source_model_name",
        "query_model_name",
        "pooling_method",
        "char_count",
        "word_count",
        "preview",
    ]
    writer = csv.DictWriter(out, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(manifest)
    return out.getvalue().encode("utf-8")


def _results_jsonl_bytes(results: list[QueryResult]) -> bytes:
    lines = []
    for r in results:
        lines.append(
            json.dumps(
                {
                    "rank": r.rank,
                    "score": r.score,
                    "pooled_id": r.pooled_id,
                    "group_key": r.group_key,
                    "source_markdown_file": r.source_markdown_file,
                    "member_count": r.member_count,
                    "embedding_dim": r.embedding_dim,
                    "source_model_name": r.source_model_name,
                    "query_model_name": r.query_model_name,
                    "pooling_method": r.pooling_method,
                    "char_count": r.char_count,
                    "word_count": r.word_count,
                    "preview": r.preview,
                    "text": r.text,
                },
                ensure_ascii=False,
            )
        )
    return ("\n".join(lines) + "\n").encode("utf-8")


def _results_npz_bytes(results: list[QueryResult], summary: dict) -> bytes:
    ids = np.asarray([r.pooled_id for r in results], dtype=object)
    group_keys = np.asarray([r.group_key for r in results], dtype=object)
    source_files = np.asarray([r.source_markdown_file for r in results], dtype=object)
    scores = np.asarray([r.score for r in results], dtype=np.float32)
    ranks = np.asarray([r.rank for r in results], dtype=np.int32)
    texts = np.asarray([r.text for r in results], dtype=object)

    bio = BytesIO()
    np.savez_compressed(
        bio,
        pooled_ids=ids,
        group_keys=group_keys,
        source_markdown_files=source_files,
        scores=scores,
        ranks=ranks,
        texts=texts,
        query_text=np.asarray([summary.get("query_text", "")], dtype=object),
        query_model_name=np.asarray([summary.get("query_model_name", "")], dtype=object),
    )
    return bio.getvalue()


def build_output_zip_bytes(
    manifest: list[dict],
    results: list[QueryResult],
    summary: dict,
    pooled_zip_name: str,
    markdown_zip_name: str,
    formats: list[str],
    include_matched_chunks: bool = True,
) -> bytes:
    stem = _safe_stem(Path(pooled_zip_name).stem + "_querying")
    bio = BytesIO()
    prompt_text = build_llm_prompt(summary.get("query_text", ""), results)

    with zipfile.ZipFile(bio, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("querying_summary.json", json.dumps(summary, ensure_ascii=False, indent=2))
        zf.writestr("querying_results.csv", _results_csv_bytes(manifest))
        zf.writestr("llm_prompt.txt", prompt_text.encode("utf-8"))
        zf.writestr(
            "inputs_used.json",
            json.dumps(
                {
                    "pooled_zip_name": pooled_zip_name,
                    "markdown_zip_name": markdown_zip_name,
                    "result_count": len(results),
                },
                ensure_ascii=False,
                indent=2,
            ),
        )

        if "jsonl" in formats:
            zf.writestr(f"{stem}_results.jsonl", _results_jsonl_bytes(results))
        if "npz" in formats:
            zf.writestr(f"{stem}_results.npz", _results_npz_bytes(results, summary))
        if include_matched_chunks:
            for r in results:
                if r.text and r.source_markdown_file:
                    zf.writestr(f"matched_chunks/{r.source_markdown_file}", r.text.encode("utf-8"))
    return bio.getvalue()
