database.py

import os
import re
from datetime import datetime, timedelta, timezone
from typing import List, Dict

import psycopg2
from psycopg2.extras import Json

# ── Configuration ────────────────────────────────────────────────────────────

# Pull DATABASE_URL from the environment, falling back to a local .env file
# so the same Postgres credentials work for both this scraper and the
# Next.js / Drizzle side without having to re-export the variable.
try:
    from dotenv import load_dotenv

    _here = os.path.dirname(os.path.abspath(__file__))
    for _candidate in (
        os.path.join(_here, ".env"),
        os.path.join(_here, "components", ".env"),
    ):
        if os.path.exists(_candidate):
            load_dotenv(_candidate, override=False)
except ImportError:
    pass

DATABASE_URL = os.environ.get("DATABASE_URL")
if not DATABASE_URL:
    raise RuntimeError(
        "DATABASE_URL is not set. Define it in your environment or in a "
        ".env file at the project root (or in components/.env)."
    )


def _connect():
    """Open a fresh Postgres connection.

    We open/close per call instead of holding a global. With Neon's pooler
    endpoint (the one in DATABASE_URL) each `connect` is effectively a
    cheap checkout from their server-side pool, so this stays simple and
    plays nicely with FastAPI's threadpool — no shared cursors between
    request handlers.
    """
    return psycopg2.connect(DATABASE_URL)


# ── Schema bootstrap ─────────────────────────────────────────────────────────


def init_database():
    """Initialize the database with required tables.

    Idempotent: every statement uses IF NOT EXISTS, so it's safe to run on
    every import and it won't fight with Drizzle's migrations if they've
    already created the table.
    """
    conn = _connect()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS scraped_jobs (
                    id              SERIAL PRIMARY KEY,
                    title           TEXT,
                    company         TEXT,
                    location        TEXT,
                    link            TEXT UNIQUE,
                    posted          TEXT,
                    posted_datetime TIMESTAMP,
                    programs        JSONB DEFAULT '[]'::jsonb,
                    schools         JSONB DEFAULT '[]'::jsonb,
                    scraped_at      TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    keywords        TEXT,
                    freshness       TEXT,
                    work_types      JSONB DEFAULT '[]'::jsonb,
                    job_types       JSONB DEFAULT '[]'::jsonb,
                    company_rating  REAL,
                    contact_details JSONB DEFAULT '[]'::jsonb
                )
                """
            )
            # Migrate existing tables that don't yet have company_rating or contact_details
            cur.execute(
                """
                ALTER TABLE scraped_jobs
                    ADD COLUMN IF NOT EXISTS company_rating REAL
                """
            )
            cur.execute(
                """
                ALTER TABLE scraped_jobs
                    ADD COLUMN IF NOT EXISTS contact_details JSONB DEFAULT '[]'::jsonb
                """
            )
            cur.execute(
                "CREATE INDEX IF NOT EXISTS idx_scraped_at "
                "ON scraped_jobs(scraped_at)"
            )
            cur.execute(
                "CREATE INDEX IF NOT EXISTS idx_link ON scraped_jobs(link)"
            )
        conn.commit()
    finally:
        conn.close()


# ── Internal helpers ─────────────────────────────────────────────────────────


def _coerce_posted_datetime(posted: str, raw_posted_dt, now_dt: datetime) -> datetime:
    """Normalize the scraper's posted_datetime into a real `datetime`.

    The scraper sometimes only gives us a relative label like "3 hours
    ago" or a bare date string. SQLite stored these as TEXT, but Postgres
    has a real TIMESTAMP column, so we resolve everything to a `datetime`
    here before insertion. Falls back to `now_dt` when neither source is
    usable. Behavior matches the original SQLite version.
    """
    needs_recompute = (
        not raw_posted_dt
        or (isinstance(raw_posted_dt, str) and len(raw_posted_dt) <= 10)
    )

    if posted and needs_recompute:
        posted_lower = posted.lower()
        match = re.search(r"(\d+)", posted_lower)
        if match:
            num = int(match.group(1))
            if "minute" in posted_lower:
                return now_dt - timedelta(minutes=num)
            if "hour" in posted_lower:
                return now_dt - timedelta(hours=num)
            if "day" in posted_lower:
                return now_dt - timedelta(days=num)
            if "week" in posted_lower:
                return now_dt - timedelta(weeks=num)
            if "month" in posted_lower:
                return now_dt - timedelta(days=num * 30)
            return now_dt

    if isinstance(raw_posted_dt, datetime):
        return raw_posted_dt
    if isinstance(raw_posted_dt, str):
        try:
            return datetime.fromisoformat(raw_posted_dt)
        except ValueError:
            pass

    return now_dt


def _row_to_job(row) -> Dict:
    """Translate a SELECT row tuple into the dict shape main.py / the UI expect.

    Notes:
      * `programs` / `schools` come back as native Python lists because
        psycopg2 auto-parses JSONB — no `json.loads` needed.
      * `posted_datetime` / `scraped_at` come back as `datetime` objects;
        we ISO-format them so the result survives `json.dumps` in FastAPI
        without extra encoders.
    """
    return {
        "title": row[0],
        "company": row[1],
        "location": row[2],
        "link": row[3],
        "posted": row[4],
        "posted_datetime": row[5].replace(tzinfo=timezone.utc).isoformat() if row[5] else None,
        "programs": row[6] if row[6] is not None else [],
        "schools": row[7] if row[7] is not None else [],
        "scraped_at": row[8].replace(tzinfo=timezone.utc).isoformat() if row[8] else None,
        "keywords": row[9],
        "company_rating": row[10],  # float or None
        "contact_details": row[11] if len(row) > 11 and row[11] is not None else [],
    }


# ── Public API (signatures preserved from the SQLite version) ────────────────


def save_scraped_jobs(jobs: List[Dict], search_params: Dict):
    """Save scraped jobs to database."""
    conn = _connect()
    saved_count = 0
    try:
        with conn.cursor() as cur:
            now_dt = datetime.now(timezone.utc)

            for job in jobs:
                posted = job.get("posted", "")
                posted_dt = _coerce_posted_datetime(
                    posted, job.get("posted_datetime"), now_dt
                )

                cur.execute(
                    """
                    INSERT INTO scraped_jobs (
                        title, company, location, link, posted,
                        posted_datetime, programs, schools, keywords,
                        freshness, work_types, job_types, scraped_at, contact_details
                    )
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (link) DO UPDATE SET
                        title           = EXCLUDED.title,
                        company         = EXCLUDED.company,
                        location        = EXCLUDED.location,
                        posted          = EXCLUDED.posted,
                        posted_datetime = EXCLUDED.posted_datetime,
                        programs        = EXCLUDED.programs,
                        schools         = EXCLUDED.schools,
                        keywords        = EXCLUDED.keywords,
                        freshness       = EXCLUDED.freshness,
                        work_types      = EXCLUDED.work_types,
                        job_types       = EXCLUDED.job_types,
                        scraped_at      = EXCLUDED.scraped_at,
                        contact_details = EXCLUDED.contact_details
                    """,
                    (
                        job.get("title"),
                        job.get("company"),
                        job.get("location"),
                        job.get("link"),
                        posted,
                        posted_dt,
                        Json(job.get("programs", []) or []),
                        Json(job.get("schools", []) or []),
                        search_params.get("keywords"),
                        search_params.get("freshness"),
                        Json(search_params.get("work_types", []) or []),
                        Json(search_params.get("job_types", []) or []),
                        now_dt,
                        Json(job.get("contact_details", []) or []),
                    ),
                )
                saved_count += 1
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()
    return saved_count


def get_jobs_by_timeframe(hours_ago: int) -> List[Dict]:
    """Get jobs scraped within the specified timeframe."""
    conn = _connect()
    try:
        with conn.cursor() as cur:
            time_threshold = datetime.now(timezone.utc) - timedelta(hours=hours_ago)
            cur.execute(
                """
                SELECT title, company, location, link, posted, posted_datetime,
                       programs, schools, scraped_at, keywords, company_rating, contact_details
                FROM scraped_jobs
                WHERE posted_datetime >= %s
                ORDER BY posted_datetime DESC
                """,
                (time_threshold,),
            )
            rows = cur.fetchall()
    finally:
        conn.close()
    return [_row_to_job(row) for row in rows]


def get_jobs_in_timeframe(start_hours: int, end_hours: int) -> List[Dict]:
    """Get jobs scraped between start_hours and end_hours ago."""
    conn = _connect()
    try:
        with conn.cursor() as cur:
            now = datetime.now(timezone.utc)
            start_time = now - timedelta(hours=start_hours)
            end_time = now - timedelta(hours=end_hours)
            cur.execute(
                """
                SELECT title, company, location, link, posted, posted_datetime,
                       programs, schools, scraped_at, keywords, company_rating, contact_details
                FROM scraped_jobs
                WHERE posted_datetime >= %s AND posted_datetime <= %s
                ORDER BY posted_datetime DESC
                """,
                (start_time, end_time),
            )
            rows = cur.fetchall()
    finally:
        conn.close()
    return [_row_to_job(row) for row in rows]


def cleanup_old_jobs(days_old: int = 7):
    """Remove jobs older than specified days."""
    conn = _connect()
    try:
        with conn.cursor() as cur:
            cutoff_time = datetime.now(timezone.utc) - timedelta(days=days_old)
            cur.execute(
                "DELETE FROM scraped_jobs WHERE posted_datetime < %s",
                (cutoff_time,),
            )
            deleted_count = cur.rowcount
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()
    return deleted_count


def get_binned_jobs() -> Dict[str, Dict]:
    """Fetch jobs from the last 30 days in one query and bin them by timeframe."""
    conn = _connect()
    try:
        with conn.cursor() as cur:
            now = datetime.now(timezone.utc)
            start_time = now - timedelta(hours=720)
            cur.execute(
                """
                SELECT title, company, location, link, posted, posted_datetime,
                       programs, schools, scraped_at, keywords, company_rating, contact_details
                FROM scraped_jobs
                WHERE posted_datetime >= %s
                ORDER BY posted_datetime DESC
                """,
                (start_time,),
            )
            rows = cur.fetchall()

            t_1, t_2, t_5, t_24, t_168, t_720 = [], [], [], [], [], []
            for row in rows:
                dt = row[5]
                job = _row_to_job(row)
                if not dt:
                    continue

                dt_aware = dt.replace(tzinfo=timezone.utc)
                hours = (now - dt_aware).total_seconds() / 3600
                if hours <= 1:
                    t_1.append(job)
                elif hours <= 2:
                    t_2.append(job)
                elif hours <= 5:
                    t_5.append(job)
                elif hours <= 24:
                    t_24.append(job)
                elif hours <= 168:
                    t_168.append(job)
                elif hours <= 720:
                    t_720.append(job)

    finally:
        conn.close()

    return {
        "1_hour": {"jobs": t_1, "count": len(t_1)},
        "2_hours": {"jobs": t_2, "count": len(t_2)},
        "5_hours": {"jobs": t_5, "count": len(t_5)},
        "24_hours": {"jobs": t_24, "count": len(t_24)},
        "1_week": {"jobs": t_168, "count": len(t_168)},
        "1_month": {"jobs": t_720, "count": len(t_720)},
    }


init_database()


def get_rated_companies() -> set:
    """Return a set of lowercased company names that already have a non-NULL rating.

    Used to skip re-fetching Serper ratings for companies we have already scored.
    """
    conn = _connect()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT DISTINCT LOWER(company)
                FROM scraped_jobs
                WHERE company_rating IS NOT NULL
                  AND company IS NOT NULL
                """
            )
            return {row[0] for row in cur.fetchall()}
    finally:
        conn.close()


def update_company_ratings(ratings: Dict[str, float]) -> int:
    """Bulk-update company_rating on all scraped_jobs rows by company name.

    Uses a single multi-row UPDATE so the number of DB round-trips equals one
    per unique company (not one per job row). Matching is case-insensitive.
    Returns the total number of rows updated.
    """
    if not ratings:
        return 0
    conn = _connect()
    updated = 0
    try:
        with conn.cursor() as cur:
            for company, rating in ratings.items():
                cur.execute(
                    """
                    UPDATE scraped_jobs
                       SET company_rating = %s
                     WHERE LOWER(company) = LOWER(%s)
                    """,
                    (float(rating), company),
                )
                updated += cur.rowcount
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()
    return updated