"""
JTBD → Funnel-Aware Keyword Seed Generator
WITH:
- Situation / Struggle / Outcome embeddings
- Noun phrase compression (keyword-like output)
- Persistent embedding cache
- Progress logging
- NO Google Ads dependency

Input:
- CSV with columns:
    situation, struggle, outcome
    (optional: job_cluster)

Output:
- jtbd_funnel_keyword_seeds.csv
"""

import os
import pickle
from pathlib import Path
from typing import List

import numpy as np
import pandas as pd
import spacy

# =============================
# CONFIG
# =============================

INPUT_CSV = "job_cluster0_hopx_jtbds.csv"
OUTPUT_CSV = "jtbd_funnel_keyword_seeds.csv"
EMBED_CACHE_PATH = Path("embedding_cache.pkl")

EMBEDDING_MODEL = "text-embedding-3-small"

SIM_HIGH = 0.65
SIM_MED = 0.45
MAX_SEEDS_LOGGED_PER_JTBD = 20

# =============================
# NLP + OPENAI
# =============================

nlp = spacy.load("en_core_web_sm")

from openai import OpenAI
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# =============================
# EMBEDDING CACHE
# =============================

def normalize_text(text: str) -> str:
    return " ".join(text.lower().strip().split()) if isinstance(text, str) else ""

if EMBED_CACHE_PATH.exists():
    with open(EMBED_CACHE_PATH, "rb") as f:
        EMBEDDING_CACHE = pickle.load(f)
    print(f"✅ Loaded embedding cache with {len(EMBEDDING_CACHE)} entries")
else:
    EMBEDDING_CACHE = {}
    print("ℹ️ No embedding cache found, starting fresh")

def embed(texts: List[str]) -> np.ndarray:
    texts = [normalize_text(t) for t in texts]

    missing = [t for t in texts if t and t not in EMBEDDING_CACHE]

    if missing:
        print(f"    🔁 Embedding {len(missing)} new texts")
        res = client.embeddings.create(
            model=EMBEDDING_MODEL,
            input=missing
        )
        for text, r in zip(missing, res.data):
            EMBEDDING_CACHE[text] = np.array(r.embedding, dtype=float)
    else:
        print("    ✅ All embeddings served from cache")

    return np.vstack([EMBEDDING_CACHE.get(t, np.zeros(1536)) for t in texts])

def cosine(a: np.ndarray, b: np.ndarray) -> float:
    denom = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / denom) if denom else 0.0

# =============================
# NOUN PHRASE COMPRESSION
# =============================

def extract_noun_phrases(text: str) -> List[str]:
    """
    Compress JTBD language into keyword-like noun phrases.
    Deterministic. No LLM.
    """
    text = normalize_text(text)
    if not text:
        return []

    doc = nlp(text)
    phrases = set()

    for chunk in doc.noun_chunks:
        phrase = normalize_text(chunk.text)

        if len(phrase.split()) < 2:
            continue
        if phrase in {"something", "anything", "everything"}:
            continue

        phrases.add(phrase)

        # Controlled expansion toward AI agent phrasing
        if "framework" in phrase and "ai" not in phrase:
            phrases.add(f"ai {phrase}")
        if "agent" in phrase and "ai" not in phrase:
            phrases.add(f"ai {phrase}")
        if "tool" in phrase and "ai" not in phrase:
            phrases.add(f"ai {phrase}")

    return list(phrases)

# =============================
# FUNNEL LABELING
# =============================

def funnel_label(sim_p, sim_o, sim_s):
    if max(sim_p, sim_o, sim_s) < SIM_MED:
        return "weak"
    if sim_p >= SIM_HIGH and sim_o < SIM_HIGH:
        return "problem"
    if sim_o >= SIM_HIGH and sim_p < SIM_HIGH:
        return "solution"
    if sim_s >= SIM_HIGH and sim_p < SIM_HIGH and sim_o < SIM_HIGH:
        return "context"
    if sim_p >= SIM_MED and sim_o >= SIM_MED:
        return "hybrid"
    return "ambiguous"

# =============================
# MAIN PIPELINE
# =============================

print("\n📥 Loading JTBD file...")
df = pd.read_csv(INPUT_CSV)

required = {"situation", "struggle", "outcome"}
missing = required - set(df.columns)
if missing:
    raise ValueError(f"Missing required columns: {missing}")

print(f"✅ Loaded {len(df)} JTBD rows")

rows = []
total = len(df)

for idx, row in df.iterrows():
    situation = row["situation"]
    struggle = row["struggle"]
    outcome = row["outcome"]
    job_cluster = row.get("job_cluster", "N/A")

    print("\n" + "=" * 70)
    print(f"[{idx+1}/{total}] JTBD {idx}")
    print(f"Job cluster: {job_cluster}")
    print(f"Situation: {str(situation)[:100]}")
    print(f"Struggle : {str(struggle)[:100]}")
    print(f"Outcome  : {str(outcome)[:100]}")

    print("  → Embedding JTBD fields")
    E_situation, E_struggle, E_outcome = embed(
        [situation, struggle, outcome]
    )

    print("  → Extracting noun-phrase seeds")
    problem_seeds  = extract_noun_phrases(struggle)
    solution_seeds = extract_noun_phrases(outcome)
    context_seeds  = extract_noun_phrases(situation)

    print(f"     Problem : {len(problem_seeds)}")
    print(f"     Solution: {len(solution_seeds)}")
    print(f"     Context : {len(context_seeds)}")

    seed_origin = {}
    for s in problem_seeds:
        seed_origin[s] = "problem"
    for s in solution_seeds:
        seed_origin.setdefault(s, "solution")
    for s in context_seeds:
        seed_origin.setdefault(s, "context")

    seeds = list(seed_origin.keys())
    print(f"  → {len(seeds)} unique compressed seeds")

    print("  → Embedding seeds")
    seed_vectors = embed(seeds)

    print("  → Scoring seeds")
    for i, (seed, vec) in enumerate(zip(seeds, seed_vectors)):
        sim_p = cosine(vec, E_struggle)
        sim_o = cosine(vec, E_outcome)
        sim_s = cosine(vec, E_situation)
        label = funnel_label(sim_p, sim_o, sim_s)

        if i < MAX_SEEDS_LOGGED_PER_JTBD:
            print(
                f"     [{label:<9}] {seed:<40} "
                f"(P:{sim_p:.2f} O:{sim_o:.2f} C:{sim_s:.2f})"
            )

        rows.append({
            "jtbd_id": idx,
            "job_cluster": job_cluster,
            "situation": situation,
            "struggle": struggle,
            "outcome": outcome,
            "seed": seed,
            "seed_origin": seed_origin[seed],
            "funnel_label": label,
            "sim_struggle": sim_p,
            "sim_outcome": sim_o,
            "sim_situation": sim_s,
        })

# =============================
# SAVE OUTPUTS
# =============================

print("\n💾 Writing output CSV...")
out_df = pd.DataFrame(rows)
out_df.to_csv(OUTPUT_CSV, index=False)
print(f"✅ Saved {len(out_df)} rows to {OUTPUT_CSV}")

print("💾 Saving embedding cache...")
with open(EMBED_CACHE_PATH, "wb") as f:
    pickle.dump(EMBEDDING_CACHE, f)
print(f"✅ Saved embedding cache ({len(EMBEDDING_CACHE)} entries)")

print("\n🎯 Pipeline complete.")