#!/usr/bin/env python3
"""Interactive QA loop from prebuilt FAISS index.

Run:
    python test/a.py
"""

from __future__ import annotations

import argparse
import importlib.util
import io
import json
import os
import sys
from pathlib import Path

# Fix Windows terminal encoding for Vietnamese
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
sys.stdin = io.TextIOWrapper(sys.stdin.buffer, encoding="utf-8")

import numpy as np
from openai import OpenAI
        
try:
    import faiss
except ImportError as exc:
    raise SystemExit(
        "faiss-cpu is required. Install with: pip install faiss-cpu"
    ) from exc


def parse_args() -> argparse.Namespace:
    here = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(
        description="Query prebuilt FAISS index in a while loop"
    )
    parser.add_argument(
        "--index", default=str(here / "index.faiss"), help="Path to FAISS index file"
    )
    parser.add_argument(
        "--chunks", default=str(here / "chunks.jsonl"), help="Path to chunk jsonl file"
    )
    parser.add_argument("--top-k", type=int, default=5, help="Top-K retrieval")
    parser.add_argument(
        "--embedding-model",
        default="text-embedding-3-small",
        help="OpenAI embedding model",
    )
    parser.add_argument("--chat-model", default="gpt-4.1-mini", help="OpenAI chat model")
    parser.add_argument(
        "--min-score", type=float, default=0.35, help="Minimum score to trust retrieval"
    )
    return parser.parse_args()


def get_api_key() -> str | None:
    backend_config = Path(__file__).resolve().parents[1] / "backend" / "config.py"
    if backend_config.exists():
        spec = importlib.util.spec_from_file_location("backend_config", backend_config)
        if spec and spec.loader:
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            key = getattr(module, "OPENAI_API_KEY", None)
            if key:
                return key
    return os.getenv("OPENAI_API_KEY")


def normalize(v: np.ndarray) -> np.ndarray:
    n = np.linalg.norm(v)
    if n == 0:
        return v.astype(np.float32)
    return (v / n).astype(np.float32)


def load_chunks(path: Path) -> list[dict]:
    rows: list[dict] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rows.append(json.loads(line))
    return rows


def embed_query(client: OpenAI, text: str, model: str) -> np.ndarray:
    response = client.embeddings.create(model=model, input=[text])
    vec = np.array(response.data[0].embedding, dtype=np.float32)
    return normalize(vec).reshape(1, -1)


def retrieve(
    index: faiss.Index,
    chunks: list[dict],
    client: OpenAI,
    query: str,
    top_k: int,
    embedding_model: str,
) -> list[dict]:
    q_vec = embed_query(client, query, embedding_model)
    scores, indices = index.search(q_vec, top_k)
    hits: list[dict] = []

    for score, idx in zip(scores[0], indices[0]):
        if idx < 0 or idx >= len(chunks):
            continue
        row = chunks[int(idx)]
        hits.append(
            {
                "chunk_id": row.get("chunk_id", int(idx)),
                "score": float(score),
                "token_start": row.get("token_start"),
                "token_end": row.get("token_end"),
                "text": row.get("text", ""),
            }
        )
    return hits


def answer_from_hits(
    client: OpenAI, query: str, hits: list[dict], chat_model: str
) -> str:
    context = "\n\n".join(f"[Chunk {h['chunk_id']}]\n{h['text']}" for h in hits)
    system_prompt = (
        "Bạn là trợ lý AI của Canifa, chuyên hỗ trợ chính sách khách hàng. "
        "Trả lời bằng tiếng Việt có dấu, tự nhiên và ngắn gọn. "
        "Sử dụng thông tin từ các đoạn context được cung cấp. "
        "Nếu câu hỏi là chào hỏi thông thường, hãy chào lại thân thiện và giới thiệu bạn có thể hỗ trợ gì. "
        "Nếu không tìm thấy thông tin liên quan, hãy nói rõ và gợi ý các chủ đề bạn có thể hỗ trợ."
    )
    user_prompt = (
        f"Câu hỏi:\n{query}\n\n"
        f"Ngữ cảnh:\n{context}\n\n"
        "Trả lời bằng tiếng Việt có dấu."
    )
    response = client.chat.completions.create(
        model=chat_model,
        temperature=0.2,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    )
    return response.choices[0].message.content or ""


def main() -> int:
    args = parse_args()
    index_path = Path(args.index)
    chunks_path = Path(args.chunks)

    if not index_path.exists():
        print(f"Missing index file: {index_path}")
        return 1
    if not chunks_path.exists():
        print(f"Missing chunks file: {chunks_path}")
        return 1

    api_key = get_api_key()
    if not api_key:
        print("Missing OPENAI_API_KEY in backend/config.py or environment")
        return 1

    print(f"Loading FAISS index from: {index_path}")
    index = faiss.read_index(str(index_path))
    chunks = load_chunks(chunks_path)
    print(f"Loaded index vectors: {index.ntotal}")
    print(f"Loaded chunks: {len(chunks)}")

    client = OpenAI(api_key=api_key)

    print("\nInteractive mode is ready.")
    print("Type exit or quit to stop.\n")

    while True:
        try:
            query = input("Ask> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nStopped.")
            break

        if not query:
            continue
        if query.lower() in {"exit", "quit"}:
            print("Stopped.")
            break

        hits = retrieve(
            index=index,
            chunks=chunks,
            client=client,
            query=query,
            top_k=args.top_k,
            embedding_model=args.embedding_model,
        )

        print("Top-K hits:")
        for h in hits:
            print(
                f"  chunk={h['chunk_id']} score={h['score']:.4f} "
                f"tokens[{h['token_start']}:{h['token_end']}]"
            )

        if not hits:
            print("Answer:")
            print("Không tìm thấy dữ liệu liên quan. Bạn thử hỏi về đổi hàng, hạng thẻ Gold/Diamond, tích điểm hoặc ưu đãi sinh nhật nhé.")
            print()
            continue

        try:
            answer = answer_from_hits(
                client=client,
                query=query,
                hits=hits,
                chat_model=args.chat_model,
            )
            print("Answer:")
            print(answer)
        except Exception as exc:
            print(f"Answer error: {exc}")

        print()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
