#!/usr/bin/env python3
"""Enrich Twenty CRM opportunities and companies from local job-search data.

Encodes the deterministic layers of the /sync-twenty enrichment pipeline:

  Opportunity salary (``amount``), in source priority order:
    1. seen_jobs.json ``salary_max_stated``   (stated in the posting at rank time)
    2. jobs.db postings salary_max            (same company+title, stated on an ATS board)
    3. jobs.db companies levels_swe_median    (levels.fyi company median - labeled estimate)

  Company fields (``domainName``, ``linkedinLink``, ``address``, ``annualRevenue``):
    - jobs.db companies row matched on normalized name, homepage domain joined
      through postings.company_domain. Stringified NaN / empty strings are absent.

Anything these layers cannot ground is emitted in the plan's ``unresolved``
section for the calling command to escalate to web research (scout subagents /
Perplexity) - this script never guesses.

Usage:
  python3 tools/twenty_enrich.py                 # plan -> job_scraper/enrichment_plan.json
  python3 tools/twenty_enrich.py --apply         # plan + apply via REST (needs TWENTY_API_KEY)
  python3 tools/twenty_enrich.py --snapshot f.json   # offline: CRM state from a JSON snapshot

Snapshot shape (what /sync-twenty exports through MCP):
  {"opportunities": [{id, url:{primaryLinkUrl}, amount, closeDate, ...}],
   "companies":     [{id, name, domainName, linkedinLink, address, ...}]}

The plan JSON is idempotent to apply: every entry names the target record id
and only fields that are currently empty.
"""
from __future__ import annotations

import argparse
import json
import re
import sqlite3
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
SEEN_JOBS = REPO / "job_scraper" / "seen_jobs.json"
JOBS_DB = REPO / "documents" / "postings" / "jobs.db"
PLAN_OUT = REPO / "job_scraper" / "enrichment_plan.json"

_SUFFIXES = re.compile(r"\b(inc|llc|ltd|corp|corporation|co|labs|lab|hq|the)\b")


def norm_company(s: str) -> str:
    """Normalized company key: lowercase, punctuation collapsed, suffixes stripped."""
    s = re.sub(r"[^a-z0-9]+", " ", (s or "").lower()).strip()
    s = _SUFFIXES.sub("", s).strip()
    return re.sub(r"\s+", " ", s)


def norm_title(s: str) -> str:
    return re.sub(r"[^a-z0-9]+", " ", (s or "").lower()).strip()


def clean_domain(raw: str | None) -> str | None:
    """Bare registrable domain or None. Rejects the stringified-NaN artifact."""
    if not raw:
        return None
    d = str(raw).strip().lower()
    for prefix in ("http://", "https://", "www."):
        d = d.removeprefix(prefix)
    d = d.rstrip("/")
    if not d or d == "nan" or "." not in d:
        return None
    return d


def _has_link(v) -> bool:
    return bool(v and v.get("primaryLinkUrl"))


def _has_addr(v) -> bool:
    return bool(v and (v.get("addressCity") or v.get("addressCountry")))


def _has_amount(v) -> bool:
    return bool(v and v.get("amountMicros"))


def load_local_sources(seen_path: Path = SEEN_JOBS, db_path: Path = JOBS_DB):
    """Ranked jobs plus the jobs.db lookup tables the plan builders consume."""
    seen = json.loads(seen_path.read_text(encoding="utf-8"))["seen"]
    ranked = [v for v in seen.values() if v.get("status") == "ranked"]

    db_companies: dict[str, dict] = {}
    domains: dict[str, str] = {}
    postings_salary: dict[tuple[str, str], int] = {}
    levels: dict[str, int] = {}
    if db_path.exists():
        con = sqlite3.connect(db_path)
        con.row_factory = sqlite3.Row
        for r in con.execute(
            "SELECT company_id, company_name, linkedin_url, country, state, apollo_city,"
            " organization_revenue, levels_swe_median_usd, levels_median_comp_usd FROM companies"
        ):
            key = norm_company(r["company_name"])
            db_companies.setdefault(key, dict(r))
            lv = r["levels_swe_median_usd"] or r["levels_median_comp_usd"]
            if lv and lv > 0:
                levels.setdefault(key, int(lv))
        for r in con.execute(
            "SELECT company_id, company_domain, COUNT(*) n FROM postings"
            " WHERE company_domain IS NOT NULL AND company_domain != ''"
            " GROUP BY company_id, company_domain ORDER BY n"
        ):
            domains[r["company_id"]] = r["company_domain"]  # last row = most frequent
        for r in con.execute(
            "SELECT company_name, title, salary_max, salary_currency FROM postings"
            " WHERE salary_max IS NOT NULL AND salary_max > 0"
        ):
            if (r["salary_currency"] or "USD").upper() in ("USD", "$", ""):
                key = (norm_company(r["company_name"]), norm_title(r["title"]))
                postings_salary.setdefault(key, int(float(r["salary_max"])))
        con.close()
    return ranked, db_companies, domains, postings_salary, levels


def build_salary_plan(ranked, opportunities, postings_salary, levels):
    """(updates, unresolved): amount patches for opportunities missing one."""
    by_url = {}
    for o in opportunities:
        u = (o.get("url") or {}).get("primaryLinkUrl")
        if u:
            by_url[u] = o
    updates, unresolved = [], []
    for v in ranked:
        o = by_url.get(v.get("url"))
        if not o or _has_amount(o.get("amount")):
            continue
        amount = source = None
        if v.get("salary_max_stated"):
            amount, source = int(v["salary_max_stated"]), "seen_jobs.salary_max_stated"
        else:
            key = (norm_company(v.get("company", "")), norm_title(v.get("title", "")))
            if key in postings_salary:
                amount, source = postings_salary[key], "jobs.db postings"
            elif norm_company(v.get("company", "")) in levels:
                amount = levels[norm_company(v.get("company", ""))]
                source = "levels.fyi company SWE median (estimate)"
        if amount:
            updates.append({
                "id": o["id"],
                "source": source,
                "patch": {"amount": {"amountMicros": amount * 1_000_000, "currencyCode": "USD"}},
            })
        else:
            unresolved.append({"id": o["id"], "url": v["url"],
                               "company": v.get("company"), "title": v.get("title")})
    return updates, unresolved


def build_company_plan(companies, db_companies, domains):
    """(updates, unresolved): fill ONLY empty company fields from jobs.db."""
    updates, unresolved = [], []
    for c in companies:
        key = norm_company(c.get("name", ""))
        m = db_companies.get(key) or {}
        patch = {}
        if not _has_link(c.get("domainName")):
            dom = clean_domain(domains.get(m.get("company_id", "")))
            if dom:
                patch["domainName"] = {"primaryLinkUrl": f"https://{dom}",
                                       "primaryLinkLabel": c["name"]}
        if not _has_link(c.get("linkedinLink")):
            li = (m.get("linkedin_url") or "").strip()
            if li:
                patch["linkedinLink"] = {"primaryLinkUrl": li.replace("http://", "https://"),
                                         "primaryLinkLabel": c["name"]}
        if not _has_addr(c.get("address")):
            addr = {k: v for k, v in [("addressCity", m.get("apollo_city")),
                                      ("addressState", m.get("state")),
                                      ("addressCountry", m.get("country"))] if v}
            if addr:
                patch["address"] = addr
        if patch:
            updates.append({"id": c["id"], "name": c["name"], "patch": patch})
        # Fields still empty after the local layer -> web-research escalation
        missing = []
        if not _has_link(c.get("domainName")) and "domainName" not in patch:
            missing.append("domain")
        if not _has_link(c.get("linkedinLink")) and "linkedinLink" not in patch:
            missing.append("linkedin")
        if not _has_addr(c.get("address")) and "address" not in patch:
            missing.append("address")
        if missing:
            unresolved.append({"id": c["id"], "name": c["name"], "missing": missing})
    return updates, unresolved


def build_name_cleanup_plan(opportunities):
    """Strip the ' — <Company>' suffix from opportunity names.

    Runs LAST, and only for opportunities whose company relation is already
    set: the suffix is the human-readable fallback while the record is
    unlinked, and becomes redundant once ``companyId`` points at the Company.
    An unlinked opportunity keeps its suffixed name so it stays identifiable.
    """
    updates = []
    for o in opportunities:
        name = o.get("name") or ""
        if o.get("companyId") and " — " in name:
            role = name.split(" — ")[0].strip()
            if role and role != name:
                updates.append({"id": o["id"], "patch": {"name": role}})
    return updates


def coverage(opportunities, companies) -> dict:
    n, m = len(companies), len(opportunities)
    return {
        "companies": {
            "total": n,
            "domain": sum(1 for c in companies if _has_link(c.get("domainName"))),
            "linkedin": sum(1 for c in companies if _has_link(c.get("linkedinLink"))),
            "address": sum(1 for c in companies if _has_addr(c.get("address"))),
        },
        "opportunities": {
            "total": m,
            "salary": sum(1 for o in opportunities if _has_amount(o.get("amount"))),
            "company_link": sum(1 for o in opportunities if o.get("companyId")),
            "score": sum(1 for o in opportunities if o.get("score") is not None),
        },
    }


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--snapshot", type=Path,
                    help="CRM state JSON (offline mode; else fetched via REST)")
    ap.add_argument("--apply", action="store_true",
                    help="apply the plan via Twenty REST (needs TWENTY_API_KEY)")
    ap.add_argument("--out", type=Path, default=PLAN_OUT, help="plan output path")
    args = ap.parse_args(argv)

    ranked, db_companies, domains, postings_salary, levels = load_local_sources()

    client = None
    if args.snapshot:
        snap = json.loads(args.snapshot.read_text(encoding="utf-8"))
        opportunities, companies = snap["opportunities"], snap["companies"]
    else:
        from twenty_client import TwentyClient  # noqa: deferred import keeps offline mode key-free
        client = TwentyClient()
        opportunities = client.list_all("opportunities")
        companies = client.list_all("companies")

    sal_updates, sal_unresolved = build_salary_plan(ranked, opportunities, postings_salary, levels)
    comp_updates, comp_unresolved = build_company_plan(companies, db_companies, domains)
    name_updates = build_name_cleanup_plan(opportunities)

    plan = {
        "coverage_before": coverage(opportunities, companies),
        "opportunity_updates": sal_updates,
        "company_updates": comp_updates,
        "name_cleanup": name_updates,
        "unresolved": {
            "opportunities_without_salary": sal_unresolved,
            "companies_missing_fields": comp_unresolved,
        },
    }

    if args.apply:
        if client is None:
            from twenty_client import TwentyClient
            client = TwentyClient()
        for u in sal_updates:
            client.update_one("opportunities", u["id"], u["patch"])
        for u in comp_updates:
            client.update_one("companies", u["id"], u["patch"])
        # Name cleanup LAST: it must observe the post-link state.
        for u in name_updates:
            client.update_one("opportunities", u["id"], u["patch"])
        plan["applied"] = {"opportunities": len(sal_updates), "companies": len(comp_updates),
                           "names_cleaned": len(name_updates)}

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(plan, indent=2), encoding="utf-8")
    cov = plan["coverage_before"]
    print(f"plan -> {args.out}")
    print(f"opportunity salary updates: {len(sal_updates)} "
          f"(coverage {cov['opportunities']['salary']}/{cov['opportunities']['total']} before)")
    print(f"company field updates: {len(comp_updates)} · name cleanups: {len(name_updates)}")
    print(f"unresolved -> web research: {len(sal_unresolved)} postings, "
          f"{len(comp_unresolved)} companies")
    if args.apply:
        print("applied via REST.")
    return 0


if __name__ == "__main__":
    sys.exit(main())