"""Data loading, feature extraction, and optional embeddings for Superpitch."""
from __future__ import annotations

import hashlib
import json
import os
import re
import subprocess
import sys
from functools import lru_cache
from typing import List

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

from superpitch.app_paths import CACHE_DIR

# Compiled once — used by both feature extraction and moonshot detection
_RE_ELITE_EMPLOYER = re.compile(
    r'(ex-?|former|previously\s+at|spent\s+\d+\s+years?\s+at|came\s+from|worked\s+at|at\s+a?\s*)'
    r'(google|meta|facebook|stripe|goldman|mckinsey|apple|amazon|microsoft|netflix|airbnb|uber|palantir|jpmorgan|morgan\s*stanley)'
)
_RE_ELITE_EDU    = re.compile(r'(phd|ph\.d|mba|stanford|mit|harvard|yale|princeton|wharton|oxford|cambridge)')
_RE_PRIOR_EXIT   = re.compile(r'previous.{0,20}exit|sold.{0,20}company|acquired')


SOURCE_TRAINING_WEIGHTS = {
    "feedback": 1.35,
    "real": 1.2,
    "contrib": 1.1,
    "synthetic": 1.0,
    "scraped": 0.65,
    "unknown": 0.9,
}

SOURCE_RETRIEVAL_WEIGHTS = {
    "feedback": 1.25,
    "real": 1.15,
    "contrib": 1.05,
    "synthetic": 0.95,
    "scraped": 0.7,
    "unknown": 0.8,
}

CASE_FEATURE_NAMES = [
    "case_neighbor_coverage",
    "case_avg_score",
    "case_approve_rate",
    "case_disagreement",
    "case_support_quality",
    "case_best_similarity",
    "case_avg_similarity",
    "case_real_rate",
    "case_scraped_rate",
    "case_conflict_flag",
]

DEFAULT_CASE_FEATURE_ROW = np.array(
    [0.0, 5.0, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    dtype=np.float64,
)


def infer_source_from_path(path: str) -> str:
    name = os.path.basename(path).lower()
    if not name:
        return "unknown"
    if "feedback" in name:
        return "feedback"
    if "real" in name:
        return "real"
    if "scraped" in name:
        return "scraped"
    if re.search(r'raw_v\d', name):
        return "synthetic"
    return "unknown"


def normalize_example_source(ex: dict, fallback_path: str = "") -> str:
    source = str(ex.get("source", "")).strip().lower()
    if source in SOURCE_TRAINING_WEIGHTS:
        return source

    for candidate in [str(ex.get("source_file", "")), fallback_path]:
        inferred = infer_source_from_path(candidate)
        if inferred != "unknown":
            return inferred

    reason = str(ex.get("key_reason", "")).strip().lower()
    if reason.startswith("auto-scored from"):
        return "scraped"
    if reason.startswith("user scored"):
        return "feedback"
    return "unknown"


def is_generic_reason(reason: str) -> bool:
    text = reason.strip().lower()
    if not text:
        return True
    generic_prefixes = (
        "auto-scored from",
        "user scored ",
    )
    return text.startswith(generic_prefixes)


def example_training_weight(ex: dict) -> float:
    source = normalize_example_source(ex)
    weight = SOURCE_TRAINING_WEIGHTS.get(source, SOURCE_TRAINING_WEIGHTS["unknown"])

    if ex.get("decision") == "approve":
        weight *= 1.15

    if source == "scraped" and is_generic_reason(ex.get("key_reason", "")):
        weight *= 0.85

    inp = ex.get("input", {})
    missing_core = sum(1 for field in ["one_liner", "traction", "team", "market"] if not inp.get(field))
    if missing_core >= 2:
        weight *= 0.9

    return round(max(weight, 0.35), 3)


def example_retrieval_weight(ex: dict) -> float:
    source = normalize_example_source(ex)
    weight = SOURCE_RETRIEVAL_WEIGHTS.get(source, SOURCE_RETRIEVAL_WEIGHTS["unknown"])
    if is_generic_reason(ex.get("key_reason", "")):
        weight *= 0.7
    return round(max(weight, 0.25), 3)


def build_tfidf_index(texts: List[str], min_df: int = 2, max_features: int = 6000) -> dict | None:
    normalized = [(text or "").strip() or "__empty__" for text in texts]
    if not normalized:
        return None

    configs = [
        {"min_df": min_df, "stop_words": "english"},
        {"min_df": 1, "stop_words": "english"},
        {"min_df": 1, "stop_words": None},
    ]
    for config in configs:
        vectorizer = TfidfVectorizer(
            ngram_range=(1, 2),
            max_features=max_features,
            **config,
        )
        try:
            matrix = vectorizer.fit_transform(normalized).astype(np.float32)
            return {"vectorizer": vectorizer, "matrix": matrix}
        except ValueError:
            continue
    return None


def build_case_feature_index(texts: List[str], examples: List[dict], min_df: int = 2) -> dict | None:
    tfidf_index = build_tfidf_index(texts, min_df=min_df, max_features=6000)
    if tfidf_index is None:
        return None

    return {
        **tfidf_index,
        "scores": np.array([float(ex.get("overall_score", ex.get("score", 5))) for ex in examples], dtype=np.float32),
        "labels": np.array([1 if ex["decision"] == "approve" else 0 for ex in examples], dtype=np.float32),
        "support_weights": np.array([example_retrieval_weight(ex) for ex in examples], dtype=np.float32),
        "real_mask": np.array([1.0 if normalize_example_source(ex) == "real" else 0.0 for ex in examples], dtype=np.float32),
        "scraped_mask": np.array([1.0 if normalize_example_source(ex) == "scraped" else 0.0 for ex in examples], dtype=np.float32),
    }


def compute_case_features(
    query_texts: List[str],
    index: dict | None,
    top_k: int = 5,
    query_matrix=None,
    self_indices=None,
) -> np.ndarray:
    if not query_texts or index is None:
        return np.tile(DEFAULT_CASE_FEATURE_ROW, (len(query_texts), 1))

    if query_matrix is None:
        query_matrix = index["vectorizer"].transform(query_texts)
    sims = cosine_similarity(query_matrix, index["matrix"])
    features = np.tile(DEFAULT_CASE_FEATURE_ROW, (len(query_texts), 1))

    for row_idx, sim_row in enumerate(sims):
        row = sim_row.copy()
        if self_indices is not None:
            row[self_indices[row_idx]] = 0.0

        ordered = row.argsort()[::-1]
        keep = [idx for idx in ordered if row[idx] > 0][:top_k]
        if not keep:
            continue

        sim = row[keep]
        support = index["support_weights"][keep]
        effective = sim * support
        denom = float(effective.sum())
        if denom <= 0:
            effective = sim
            denom = float(effective.sum())
        if denom <= 0:
            continue

        approve_rate = float(np.dot(effective, index["labels"][keep]) / denom)
        avg_score = float(np.dot(effective, index["scores"][keep]) / denom)
        support_quality = float(np.dot(effective, support) / denom)
        real_rate = float(np.dot(effective, index["real_mask"][keep]) / denom)
        scraped_rate = float(np.dot(effective, index["scraped_mask"][keep]) / denom)
        disagreement = float(2 * min(approve_rate, 1 - approve_rate))
        best_similarity = float(sim.max())
        avg_similarity = float(np.mean(sim))
        conflict_flag = 1.0 if best_similarity >= 0.25 and 0.25 < approve_rate < 0.75 else 0.0

        features[row_idx] = np.array([
            len(keep) / float(top_k),
            avg_score,
            approve_rate,
            disagreement,
            support_quality,
            best_similarity,
            avg_similarity,
            real_rate,
            scraped_rate,
            conflict_flag,
        ], dtype=np.float64)

    return features


def case_feature_dict(row: np.ndarray) -> dict:
    context = {}
    for name, value in zip(CASE_FEATURE_NAMES, row):
        key = name[len("case_"):] if name.startswith("case_") else name
        context[key] = float(value)
    return context


def _probe_python(snippet: str) -> bool:
    try:
        result = subprocess.run(
            [sys.executable, "-c", snippet],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=20,
        )
    except Exception:
        return False
    return result.returncode == 0


def _env_truthy(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() in {"1", "true", "yes", "on"}


@lru_cache(maxsize=1)
def mlx_backend_available() -> bool:
    if _env_truthy("SUPERPITCH_DISABLE_MLX"):
        return False
    return _probe_python("import mlx.core as mx; print('ok')")


@lru_cache(maxsize=1)
def sentence_transformers_backend_available() -> bool:
    if _env_truthy("SUPERPITCH_DISABLE_EMBEDDINGS"):
        return False
    return _probe_python(
        "from sentence_transformers import SentenceTransformer; "
        "SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2'); print('ok')"
    )


@lru_cache(maxsize=1)
def preferred_embedding_backend() -> str | None:
    forced = os.environ.get("SUPERPITCH_EMBEDDING_BACKEND", "").strip().lower()
    if forced in {"mlx", "sentence-transformers"}:
        if forced == "mlx" and mlx_backend_available():
            return "mlx"
        if forced == "sentence-transformers" and sentence_transformers_backend_available():
            return "sentence-transformers"
        return None

    if mlx_backend_available() and not _env_truthy("SUPERPITCH_DISABLE_EMBEDDINGS"):
        return "mlx"
    if sentence_transformers_backend_available():
        return "sentence-transformers"
    return None


@lru_cache(maxsize=1)
def embedding_backend_available() -> bool:
    return preferred_embedding_backend() is not None


def pitch_to_text(inp: dict) -> str:
    """Flatten a pitch memo dict into a single text string for the model."""
    parts = []
    for field in ["one_liner", "stage", "raising", "traction", "team",
                  "market", "business_model", "competition", "summary"]:
        if inp.get(field):
            parts.append(f"{field}: {inp[field]}")
    if inp.get("red_flags"):
        parts.append(f"red_flags: {'; '.join(inp['red_flags'])}")
    if inp.get("strengths"):
        parts.append(f"strengths: {'; '.join(inp['strengths'])}")
    return " | ".join(parts)


def _parse_dollar(match, val_groups, suffix_groups):
    """Parse a dollar amount with K/M suffix, normalize to thousands.

    $45K → 45, $2M → 2000, $500 (bare) → 0.5
    """
    raw = None
    suffix = None
    for vg, sg in zip(val_groups, suffix_groups):
        try:
            if match.group(vg):
                raw = float(match.group(vg).replace(",", ""))
                suffix = (match.group(sg) or "").lower()
                break
        except (IndexError, TypeError):
            continue
    if raw is None:
        return 0
    if suffix == "m":
        return raw * 1000  # millions → thousands
    elif suffix == "k":
        return raw  # already in thousands
    else:
        return raw / 1000  # bare number → thousands


def extract_features(text: str, inp: dict) -> np.ndarray:
    """Hand-crafted features that predict pitch quality.

    These encode what investors actually look for: revenue traction,
    retention metrics, team credentials, risk signals, unit economics.
    """
    f = []
    t = text.lower()

    # Revenue signals (handle both "$45K MRR" and "MRR is $45K" patterns)
    # Values normalized to thousands (K) so $45K=45, $2M=2000, bare $500=0.5
    f.append(1 if re.search(r'\$[\d.]+[km]?\s*(mrr|arr)|(?:mrr|arr)\s*(?:is|of|at|:)?\s*\$', t) else 0)
    mrr_match = re.search(r'(?:\$([\d,.]+)\s*(k|m)?\s*mrr|mrr\s*(?:is|of|at|:)?\s*\$([\d,.]+)\s*(k|m)?)', t)
    mrr_raw = _parse_dollar(mrr_match, [1, 3], [2, 4]) if mrr_match else 0
    f.append(mrr_raw)
    arr_match = re.search(r'(?:\$([\d,.]+)\s*(k|m)?\s*a(?:nnualized|rr)|(?:arr|annualized)\s*(?:is|of|at|:)?\s*\$([\d,.]+)\s*(k|m)?)', t)
    arr_raw = _parse_dollar(arr_match, [1, 3], [2, 4]) if arr_match else 0
    f.append(np.log1p(arr_raw))

    # Traction signals
    f.append(1 if re.search(r'(customer|client|user)s?', t) else 0)
    f.append(1 if re.search(r'\d+%.*growth', t) else 0)
    f.append(1 if re.search(r'(retention|nrr|ndr)', t) else 0)
    # Handle "125% NRR", "NRR is 125%", "net revenue retention is 125%", "125% net revenue retention"
    nrr_match = re.search(r'(?:(\d+)%\s*(?:net\s*)?(?:revenue\s*)?retention|(?:net\s*)?(?:revenue\s*)?retention\s*(?:is|of|at|:)?\s*(\d+)%|(\d+)%\s*nrr|nrr\s*(?:is|of|at|:)?\s*(\d+)%)', t)
    if nrr_match:
        nrr_val = next(g for g in nrr_match.groups() if g is not None)
        f.append(float(nrr_val))
    else:
        f.append(0)
    churn_match = re.search(r'(?:([\d.]+)%\s*(?:logo\s*)?churn|churn\s*(?:is|of|at|:)?\s*([\d.]+)%)', t)
    f.append(float((churn_match.group(1) or churn_match.group(2))) if churn_match else 0)
    f.append(1 if re.search(r'zero\s*churn|0%?\s*churn|no\s*churn', t) else 0)

    # Revenue quality
    f.append(1 if re.search(r'\$0\s*revenue|pre-revenue|no\s*revenue', t) else 0)
    growth_match = re.search(r'(?:(\d+)%\s*mo(?:m|nth).*growth|growth\s*(?:is|of|at|:)?\s*(\d+)%\s*mo)', t)
    f.append(float((growth_match.group(1) or growth_match.group(2))) if growth_match else 0)

    # Team signals (handle "at Goldman", "from Stripe", "ex-Google", "previously at McKinsey")
    f.append(1 if _RE_ELITE_EMPLOYER.search(t) else 0)
    f.append(1 if _RE_ELITE_EDU.search(t) else 0)
    f.append(1 if _RE_PRIOR_EXIT.search(t) else 0)
    f.append(1 if 'first startup' in t or 'first-time' in t or 'first company' in t else 0)
    f.append(1 if 'solo founder' in t else 0)
    f.append(1 if re.search(r'part.time|advisor.{0,10}(not|isn)', t) else 0)

    # Risk signals
    n_red = len(inp.get('red_flags', []))
    n_str = len(inp.get('strengths', []))
    f.append(n_red)
    f.append(n_str)
    f.append(n_str - n_red)

    # Market & regulatory
    f.append(1 if re.search(r'(regulatory|fda|compliance|sec\s)', t) else 0)
    f.append(1 if re.search(r'(niche|small market|narrow|small\s*addressable)', t) else 0)
    f.append(1 if re.search(r'commodit|race.to.bottom', t) else 0)

    # Unit economics (handle "82% gross margin" and "gross margins are 82%")
    gm_match = re.search(r'(?:(\d+)%\s*(?:gross\s*)?margin|(?:gross\s*)?margins?\s*(?:is|are|of|at|:)?\s*(\d+)%)', t)
    f.append(float((gm_match.group(1) or gm_match.group(2))) if gm_match else 0)

    # Gross profit / profitability signal
    f.append(1 if re.search(r'(profitable|gross\s*profit\s*positive|unit\s*economics\s*positive|contribution\s*margin\s*positive)', t) else 0)

    # Competitive signals
    f.append(1 if re.search(r'(no\s*clear|lacks?|weak)\s*(moat|differentiat|defensib)', t) else 0)
    f.append(1 if re.search(r'(point.solution|feature.not.company|bundling\s*risk)', t) else 0)

    # Revenue per customer (mrr_value / customer_count if both > 0)
    mrr_val = f[1] if len(f) > 1 else 0  # mrr_value index
    cust_match2 = re.search(r'(\d[\d,]*)\s*(?:paying\s*)?(?:customer|client)s?\b', t)
    cust_count = float(cust_match2.group(1).replace(",", "")) if cust_match2 else 0
    if mrr_val > 0 and cust_count > 0:
        rev_per_cust = mrr_val / cust_count
    else:
        rev_per_cust = 0
    f.append(rev_per_cust)

    # Raise amount — parse "$500K seed", "$2M Series A" from raising field
    raising_text = inp.get("raising", "") or ""
    raise_match = re.search(r'\$([\d,.]+)\s*(k|m|b)?', raising_text.lower())
    if raise_match:
        raw = float(raise_match.group(1).replace(",", ""))
        sfx = (raise_match.group(2) or "").lower()
        if sfx == "b":
            raise_k = raw * 1_000_000
        elif sfx == "m":
            raise_k = raw * 1000
        elif sfx == "k":
            raise_k = raw
        else:
            raise_k = raw / 1000
    else:
        raise_k = 0
    f.append(np.log1p(raise_k))

    # Stage
    f.append(1 if re.search(r'pre.seed', t) else 0)
    f.append(1 if re.search(r'\bseed\b', t) else 0)
    f.append(1 if re.search(r'series\s*a', t) else 0)

    # Text length (longer pitches tend to have more substance)
    f.append(np.log1p(len(t)))

    # has_arr — binary signal for ARR specifically (separate from has_mrr_arr)
    f.append(1 if re.search(r'\$[\d.]+[km]?\s*a(?:nnualized|rr)|(?:arr|annualized)\s*(?:is|of|at|:)?\s*\$', t) else 0)

    # YoY growth percentage — extract "200% year-over-year growth", "3x YoY"
    yoy_match = re.search(r'(\d+)%\s*(?:year.over.year|yoy|annual)\s*(?:growth|increase)', t)
    if not yoy_match:
        yoy_match = re.search(r'(\d+)x\s*(?:yoy|year.over.year|annual)\s*growth', t)
    f.append(float(yoy_match.group(1)) if yoy_match else 0)

    # Pre-revenue with high red flags: most dangerous combination
    pre_rev_flag = 1 if re.search(r'\$0\s*revenue|pre-revenue|no\s*revenue', t) else 0
    f.append(pre_rev_flag * n_red)

    return np.array(f, dtype=np.float32)


# Feature names for interpretability
FEATURE_NAMES = [
    "has_mrr_arr", "mrr_value", "arr_value",
    "has_customers", "has_growth_pct", "has_retention", "nrr_value", "churn_pct", "zero_churn",
    "pre_revenue", "mom_growth_pct",
    "elite_employer", "elite_edu", "prior_exit", "first_time", "solo_founder", "part_time_team",
    "n_red_flags", "n_strengths", "strength_minus_red",
    "regulatory", "niche_market", "commoditized",
    "gross_margin",
    "is_profitable", "weak_moat", "point_solution",
    "rev_per_customer",
    "raise_amount_k",
    "pre_seed", "seed", "series_a",
    "text_length",
    "has_arr",
    "yoy_growth_pct",
    "pre_rev_red_flags",
]


def load_examples(path: str) -> List[dict]:
    """Load raw examples and normalize source metadata when possible."""
    examples = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            ex = json.loads(line)
            if not isinstance(ex.get("input"), dict):
                ex["input"] = {"one_liner": ex.get("text", "")[:200]}
            elif ex.get("text") and not ex["input"].get("one_liner"):
                ex["input"] = {**ex["input"], "one_liner": ex["text"][:200]}
            ex.setdefault("source_file", os.path.basename(path))
            ex["source"] = normalize_example_source(ex, path)
            examples.append(ex)
    return examples


def _cache_path(data_path: str, n_texts: int | None = None) -> str:
    with open(data_path, "rb") as f:
        file_hash = hashlib.md5(f.read()).hexdigest()[:12]
    base = os.path.splitext(os.path.basename(data_path))[0]
    suffix = f"_n{n_texts}" if n_texts is not None else ""
    backend = preferred_embedding_backend() or "none"
    return str(CACHE_DIR / f"{base}_{backend}_{file_hash}{suffix}.npy")


_embedding_model_cache: dict = {}


def _get_embedding_model(backend: str):
    if backend not in _embedding_model_cache:
        if backend == "mlx":
            from mlx_embedding_models.embedding import EmbeddingModel
            _embedding_model_cache[backend] = EmbeddingModel.from_registry("minilm-l6")
        elif backend == "sentence-transformers":
            from sentence_transformers import SentenceTransformer
            _embedding_model_cache[backend] = SentenceTransformer(
                "sentence-transformers/all-MiniLM-L6-v2"
            )
    return _embedding_model_cache[backend]


def embed_texts(texts: List[str], data_path: str | None = None, n_texts: int | None = None) -> np.ndarray:
    """Encode texts with the best available embedding backend."""
    if not embedding_backend_available():
        raise RuntimeError("Embedding backend unavailable in this environment.")

    if data_path:
        cache = _cache_path(data_path, n_texts)
        if os.path.exists(cache):
            cached = np.load(cache)
            if n_texts is None or cached.shape[0] == n_texts:
                return cached

    backend = preferred_embedding_backend()
    model = _get_embedding_model(backend)
    if backend == "mlx":
        embeddings = model.encode(texts)
    elif backend == "sentence-transformers":
        embeddings = model.encode(texts, convert_to_numpy=True, normalize_embeddings=False)
    else:
        raise RuntimeError("No supported embedding backend is available.")

    if data_path:
        os.makedirs(CACHE_DIR, exist_ok=True)
        np.save(cache, embeddings)

    return np.asarray(embeddings, dtype=np.float32)
