from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from qa.models import PermenparDataset


def normalize_kbli(val):
    """Ambil digit saja dari KBLI."""
    return "".join(filter(str.isdigit, str(val)))


class SemanticSearchAI:
    def __init__(self):
        self.vectorizer = TfidfVectorizer(
            ngram_range=(1, 2),
            max_features=5000
        )

    def _build_index(self, rows):
        documents = []
        records = []

        for row in rows:
            teks = row.teks or ""
            keywords = row.kata_kunci or ""

            documents.append(
                f"{row.kategori or ''} "
                f"{row.sub_kategori or ''} "
                f"{teks} "
                f"{keywords}"
            )
            records.append(row)

        matrix = self.vectorizer.fit_transform(documents)
        return matrix, records

    def detect_kategori(self, query):
        q = query.lower()

        if any(k in q for k in ["wajib", "kewajiban", "harus"]):
            return "kewajiban"

        if any(k in q for k in ["izin", "perizinan", "persyaratan", "syarat"]):
            return "perizinan"

        if any(k in q for k in ["sdm", "tenaga kerja"]):
            return "sdm"

        if any(k in q for k in ["sarana", "prasarana"]):
            return "sarana"

        if any(k in q for k in ["manajemen"]):
            return "manajemen"

        if any(k in q for k in ["terangkan", "jelaskan", "tentang", "apa itu"]):
            return "definisi"

        return None

    def search(self, query, kbli=None, top_n=5):
        # =========================
        # AMBIL DATA AKTIF
        # =========================
        base_qs = PermenparDataset.objects.filter(status_aktif=True)

        # =========================
        # FILTER KBLI (AMAN)
        # =========================
        rows = list(base_qs)

        if kbli:
            kbli_norm = normalize_kbli(kbli)
            rows = [
                r for r in rows
                if normalize_kbli(r.kbli) == kbli_norm
            ]

            if not rows:
                return []

        # =========================
        # DETEKSI KATEGORI
        # =========================
        kategori = self.detect_kategori(query)

        if kategori:
            rows_kat = [
                r for r in rows
                if kategori.lower() in (r.kategori or "").lower()
            ]
            if rows_kat:
                rows = rows_kat

        if not rows:
            return []

        # =========================
        # MODE UMUM KBLI (HANYA JIKA MURNI UMUM)
        # =========================
        query_lower = query.lower()
        is_umum = (
            kbli
            and kategori == "definisi"
            and not any(k in query_lower for k in ["wajib", "izin", "syarat"])
        )

        if is_umum:
            rows = sorted(
                rows,
                key=lambda x: (x.pasal or 0, x.sub_pasal or 0)
            )[:top_n]

            return [{"score": 1.0, "data": r} for r in rows]

        # =========================
        # MODE SEMANTIC TF-IDF
        # =========================
        matrix, records = self._build_index(rows)
        query_vec = self.vectorizer.transform([query])
        similarity = cosine_similarity(query_vec, matrix)[0]

        ranked = sorted(
            enumerate(similarity),
            key=lambda x: x[1],
            reverse=True
        )

        results = []
        for idx, score in ranked[:top_n]:
            if score > 0.01:
                results.append({
                    "score": round(float(score), 3),
                    "data": records[idx]
                })

        return results
