import pandas as pd
import numpy as np
import re
import sys
from pathlib import Path
from sklearn.model_selection import train_test_split
from Sastrawi.Stemmer.StemmerFactory import StemmerFactory

factory = StemmerFactory()
stemmer = factory.create_stemmer()

BASE_DIR = Path(__file__).parent
DATA_DIR = BASE_DIR
OUTPUT_DIR = DATA_DIR / "processed"
OUTPUT_DIR.mkdir(exist_ok=True)
FILE_INTERNAL = DATA_DIR / "ts_internal_mif.xlsx"
SEPARATOR = ";"
ENCODING = "utf-8-sig"
TARGET_CLASSES = ["Programmer", "Data Analyst", "Wirausaha Informatika", "Non-IT"]
KLASIFIKASI_MAP = {
    "Programmer": "Programmer",
    "Data Analyst": "Data Analyst",
    "Wirausaha IT": "Wirausaha Informatika",
    "Wirausaha": "Wirausaha Informatika",
    "Non-IT": "Non-IT",
    "Infokom": None, "Pelajar": None, "Tidak Bekerja": None, "TIdak diketahui": None
}
KEYWORD_RULES = {
    "Programmer": ["programmer", "developer", "engineer", "fullstack", "backend", "frontend", "mobile", "android", "ios", "software", "web dev", "coding", "it staff", "teknisi", "sistem informasi", "application", "network", "devops", "qa", "tester", "ui", "ux", "swe"],
    "Data Analyst": ["data analyst", "analis data", "data science", "business analyst", "research", "statistik", "bi analyst", "reporting", "database", "sql", "etl", "data engineer", "big data", "analyst", "data mining", "machine learning", "data visual", "power bi", "tableau", "looker", "business intelligence", "bi developer", "data warehouse"],
    "Wirausaha Informatika": ["founder", "owner", "ceo", "wiraswasta", "startup", "freelance", "freelancer", "wirausaha", "bisnis", "usaha mandiri", "konsultan", "co founder", "entrepreneur", "self employed", "owner toko", "usaha", "dagang online", "tokopedia", "shopee", "dropship", "reseller"]
}
COMPANY_STOPWORDS = {
    "pt", "cv", "ud", "tbk", "persero", "corp", "inc", "ltd", "koperasi", "bumn", "bumd",
    "dinas", "kantor", "pemkab", "pemprov", "politeknik", "universitas", "sekolah", "sma", "smk", "sd",
    "bank", "bpr", "rs", "rumah sakit", "klinik", "apotek", "hotel", "restoran", "cafe", "toko", "konter",
    "foundation", "yayasan", "perkumpulan", "organisasi", "agency", "studio", "consulting", "group", "holding"
}

def load_data_file(file_path: str) -> pd.DataFrame:
    ext = Path(file_path).suffix.lower()
    if ext == '.csv':
        try: return pd.read_csv(file_path, sep=SEPARATOR, encoding=ENCODING, dtype=str, on_bad_lines='skip', engine='python')
        except UnicodeDecodeError: return pd.read_csv(file_path, sep=SEPARATOR, encoding='latin1', dtype=str, on_bad_lines='skip', engine='python')
    elif ext == '.xlsx':
        return pd.read_excel(file_path, dtype=str)
    else:
        raise ValueError(f"Format tidak didukung: {ext}")

def find_col(df: pd.DataFrame, keywords: list) -> str | None:
    for col in df.columns:
        col_clean = str(col).strip().lower()
        if any(k.strip().lower() in col_clean for k in keywords): return col
    return None

def clean_text(text: str) -> str:
    if pd.isna(text) or str(text).strip().lower() in ["nan", "none", "null", "-", "0", "", "tidak diisi", "tidak diketahui"]: return ""
    text = str(text).strip().lower()
    text = re.sub(r'^\d+\s*[-:/]\s*', '', text)
    text = text.replace('-', ' ').replace('/', ' ').replace('_', ' ').replace(',', ' ')
    text = re.sub(r'[^\w\s]', '', text)
    text = re.sub(r'\d+', '', text)
    tokens = [w for w in text.split() if w not in COMPANY_STOPWORDS and len(w) >= 3]
    return " ".join([stemmer.stem(w) for w in tokens]) 

def is_likely_name(text: str, full_name: str) -> bool:
    if not text or not full_name: return False
    text_clean = re.sub(r'[^\w\s]', '', text.lower())
    name_clean = re.sub(r'[^\w\s]', '', str(full_name).lower())
    text_parts = set(text_clean.split())
    name_parts = set(name_clean.split())
    
    if len(text_parts) == 0: return False
    # Logika yang lebih ketat: Cek apakah input hanya berisi komponen nama (indikasi bocor nama)
    if len(text_parts) <= 3 and text_parts.issubset(name_parts):
        return True
    return False

def classify_rule_based(text: str) -> str:
    if not text or len(text) < 3: return "Non-IT"
    for profile in ["Programmer", "Data Analyst", "Wirausaha Informatika"]:
        if any(kw in text for kw in KEYWORD_RULES[profile]): return profile
    return "Non-IT"

def main():
    if not FILE_INTERNAL.exists():
        print(f"Berkas tidak ditemukan: {FILE_INTERNAL}"); sys.exit(1)
    print("[1/4] Memuat & Membersihkan Data (Strict Cleaning)...")
    df = load_data_file(str(FILE_INTERNAL))
    df.columns = df.columns.str.strip()
    col_nim = find_col(df, ["nim"])
    col_nama = find_col(df, ["nama", "lengkap"])
    col_jab_lama = find_col(df, ["jabatan"])
    col_jab_baru = find_col(df, ["jabatan_terupdate"])
    col_klasifikasi = find_col(df, ["klasifikasi"])
    col_status = find_col(df, ["status", "kerja"])
    
    if not all([col_nim, col_nama, col_jab_lama]):
        print(" Kolom esensial tidak ditemukan"); print(df.columns.tolist()[:10]); sys.exit(1)
        
    if col_jab_baru:
        mask_empty = df[col_jab_baru].isna() | df[col_jab_baru].astype(str).str.strip().str.lower().isin(["", "-", "0", "nan", "none", "null", "tidak diisi", "tidak diketahui"])
        df["jabatan_final"] = df[col_jab_baru].where(~mask_empty, df[col_jab_lama])
    else:
        df["jabatan_final"] = df[col_jab_lama]
        
    col_jabatan = "jabatan_final"
    if col_status:
        blacklist = ["tidah diketahui", "tidak bekerja", "pelajar", "melanjutkan pendidikan", "nan", ""]
        mask = ~df[col_status].str.lower().str.strip().isin(blacklist)
        df = df[mask].copy()
    
    # Stemming berat terjadi HANYA di sini
    df["job_text_raw"] = df[col_jabatan].apply(clean_text)
    
    if col_klasifikasi:
        fallback_map = {"Programmer": "programmer developer", "Data Analyst": "data analyst", "Wirausaha IT": "wirausaha founder", "Wirausaha": "wirausaha founder", "Non IT": "staff admin", "Infokom": "it staff teknisi", "TIdah diketahui": "", "Pelajar": "", "Tidak Bekerja": ""}
        empty_mask = df["job_text_raw"] == ""
        if empty_mask.any(): df.loc[empty_mask, "job_text_raw"] = df.loc[empty_mask, col_klasifikasi].map(fallback_map).fillna("")
    if col_nama:
        name_leak_mask = df.apply(lambda row: is_likely_name(row["job_text_raw"], row[col_nama]), axis=1)
        leaked_count = name_leak_mask.sum()
        if leaked_count > 0: print(f" Mengabaikan {leaked_count} baris yang terindikasi mengandung nama pribadi.")
        df = df[~name_leak_mask].copy()
    if col_klasifikasi:
        df["label"] = df[col_klasifikasi].str.strip().map(KLASIFIKASI_MAP)
        missing = df["label"].isna()
        if missing.any(): df.loc[missing, "label"] = df.loc[missing, "job_text_raw"].apply(classify_rule_based)
    else: df["label"] = df["job_text_raw"].apply(classify_rule_based)
    
    df = df[df["job_text_raw"].str.len() >= 3].copy()
    result = df[[col_nim, col_jabatan, "job_text_raw", "label"]].rename(columns={col_nim: "nim"}).dropna(subset=["nim", "label"]).drop_duplicates(subset=["nim"], keep="first")
    
    print(f"Memuat {len(result)} data valid (berdasarkan teks pekerjaan)\n")
    
    cols_out = ["nim", "job_text_raw", "label"]
    
    # Kita tetap simpan data uji ke CSV (jika butuh untuk sanity check statis), tapi tidak akan digunakan 
    # oleh retrain_worker secara buta lagi.
    min_count = result["label"].value_counts().min()
    if min_count < 2: train_df, test_df = train_test_split(result, test_size=0.30, random_state=42)
    else: train_df, test_df = train_test_split(result, test_size=0.30, stratify=result["label"], random_state=42)
    
    train_df[cols_out].to_csv(OUTPUT_DIR / "training_corpus_4.csv", index=False, sep=";", encoding=ENCODING)
    test_df[cols_out].to_csv(OUTPUT_DIR / "test_set_4.csv", index=False, sep=";", encoding=ENCODING)
    
    print("DATA PELATIHAN (TRAINING SET):")
    print(train_df["label"].value_counts().to_string())
    print(f"\nData berhasil disimpan ke {OUTPUT_DIR}")

if __name__ == "__main__":
    main()