import argparse
import json
import os
import re

import faiss
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer
from tqdm import tqdm


def simple_clean(text: str) -> str:
    if not isinstance(text, str):
        return ""
    text = text.replace("\u00a0", " ")
    text = re.sub(r"\s+", " ", text).strip()
    return text


def chunk_text(text: str, chunk_chars: int = 900, overlap: int = 150):
    """
    Simple char-based chunking (good enough for reviews).
    For better chunking, split by sentences and cap token length.
    """
    text = simple_clean(text)
    if len(text) <= chunk_chars:
        return [text] if text else []
    chunks = []
    i = 0
    while i < len(text):
        chunk = text[i : i + chunk_chars]
        if chunk:
            chunks.append(chunk)
        i += max(1, chunk_chars - overlap)
    return chunks


def detect_text_col(df: pd.DataFrame) -> str:
    # Heuristic: pick the longest average string column
    best_col, best_score = None, -1
    for col in df.columns:
        sample = df[col].dropna().astype(str).head(200)
        if len(sample) == 0:
            continue
        avg_len = sample.map(len).mean()
        if avg_len > best_score:
            best_score = avg_len
            best_col = col
    if best_col is None:
        raise ValueError("Could not detect a text column in the .csv file.")
    return best_col


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input_csv", required=True, help="Tripadvisor reviews .csv file")
    ap.add_argument("--out_dir", default="out")
    ap.add_argument(
        "--embedding_model", default="sentence-transformers/all-MiniLM-L6-v2"
    )
    ap.add_argument("--chunk_chars", type=int, default=900)
    ap.add_argument("--overlap", type=int, default=150)
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    df = pd.read_csv(args.input_csv, sep=",", dtype=str, on_bad_lines="skip")

    rows = df["Original"].fillna("").astype(str).tolist()

    corpus_path = os.path.join(args.out_dir, "corpus.jsonl")
    corpus = []
    doc_id = 0

    for r in tqdm(rows, desc="Chunking"):
        r = simple_clean(r)
        if len(r) < 30:
            continue
        chunks = chunk_text(r, chunk_chars=args.chunk_chars, overlap=args.overlap)
        for ch in chunks:
            if len(ch) < 30:
                continue
            corpus.append({"doc_id": doc_id, "text": ch})
            doc_id += 1

    with open(corpus_path, "w", encoding="utf-8") as f:
        for ex in corpus:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")

    # Build FAISS index
    embedder = SentenceTransformer(args.embedding_model)
    texts = [c["text"] for c in corpus]
    embs = embedder.encode(
        texts, batch_size=64, show_progress_bar=True, normalize_embeddings=True
    )
    embs = np.asarray(embs, dtype=np.float32)

    dim = embs.shape[1]
    index = faiss.IndexFlatIP(dim)  # cosine if normalized
    index.add(embs)

    faiss_path = os.path.join(args.out_dir, "faiss.index")
    faiss.write_index(index, faiss_path)

    # Store mapping doc row -> text
    mapping_path = os.path.join(args.out_dir, "docstore.jsonl")
    with open(mapping_path, "w", encoding="utf-8") as f:
        for i, c in enumerate(corpus):
            f.write(json.dumps({"faiss_id": i, **c}, ensure_ascii=False) + "\n")

    print(f"Saved:\n- {corpus_path}\n- {faiss_path}\n- {mapping_path}")


if __name__ == "__main__":
    main()