db.py

"""
Single shared data-access layer for Postgres (Supabase).

Everything that used to be scattered across agent.py and main.py (raw
psycopg2 connect/execute/commit calls, repeated in a dozen places) lives here
instead. agent.py and the Streamlit pages both import from this module so
there's one source of truth for the schema's shape and one place to fix bugs.

Uses a small connection pool (not a fresh connection per call) since
Streamlit reruns the whole script on every interaction — without pooling
you'd open a new Postgres connection on every button click.
"""

import threading
import psycopg2
from psycopg2 import pool
from psycopg2.extras import RealDictCursor
from contextlib import contextmanager

from config import get_secret

_pool = None
_pool_lock = threading.Lock()


def _get_pool():
    global _pool
    if _pool is None:
        with _pool_lock:
            if _pool is None:  # re-check inside the lock
                dsn = get_secret("DATABASE_URL")
                if not dsn:
                    raise RuntimeError(
                        "DATABASE_URL is not set. Locally: add it to your .env file. "
                        "On Streamlit Cloud: add it under your app's Settings -> Secrets."
                    )
                _pool = psycopg2.pool.SimpleConnectionPool(
                    1, 10, dsn, cursor_factory=RealDictCursor
                )
    return _pool


@contextmanager
def get_cursor(commit: bool = False):
    """
    Usage:
        with get_cursor() as cur:
            cur.execute("SELECT * FROM products WHERE id = %s", (pid,))
            row = cur.fetchone()

        with get_cursor(commit=True) as cur:
            cur.execute("UPDATE products SET ...")
    """
    conn = _get_pool().getconn()
    try:
        cur = conn.cursor()
        yield cur
        if commit:
            conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        cur.close()
        _get_pool().putconn(conn)


# ---------------------------------------------------------------------- #
# Products
# ---------------------------------------------------------------------- #

def save_product(product: dict, keyword: str = None, destination_country: str = None,
                  destination_area: str = None, run_id: str = None) -> str:
    """Inserts a product row and returns its new id."""
    import json
    with get_cursor(commit=True) as cur:
        cur.execute(
            """
            INSERT INTO products (title, raw_payload, supplier_cost, shipping_fee,
                                   optimized_price, status, approval_status, video_prompt,
                                   video_url, rejection_reason, replaces_product_id,
                                   run_keyword, run_destination_country, run_destination_area, run_id)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING id
            """,
            (
                product.get("title"),
                json.dumps(product),
                product.get("supplier_cost"),
                product.get("shipping_fee"),
                product.get("optimized_price"),
                product.get("status", "raw"),
                product.get("approval_status", "pending"),
                product.get("video_prompt"),
                product.get("video_url"),
                product.get("rejection_reason"),
                product.get("replaces_product_id"),
                keyword,
                destination_country,
                destination_area,
                run_id,
            ),
        )
        return str(cur.fetchone()["id"])


def update_product(product_id: str, **fields) -> None:
    """Generic partial update: update_product(pid, status='ready_to_deploy', approval_status='pending')"""
    if not fields:
        return
    set_clauses = [f"{k} = %s" for k in fields]
    values = list(fields.values()) + [product_id]
    with get_cursor(commit=True) as cur:
        cur.execute(f"UPDATE products SET {', '.join(set_clauses)} WHERE id = %s", values)


def merge_raw_payload(product_id: str, patch: dict) -> None:
    """Merges new keys into an existing product's raw_payload jsonb instead of
    overwriting the column. save_product() only snapshots raw_payload at
    pricing time — creative fields (store_description, social_caption) get
    added to the product dict afterwards by asset_generator_node and need to
    be written back explicitly, or the dashboard silently shows stale data."""
    import json
    with get_cursor(commit=True) as cur:
        cur.execute(
            "UPDATE products SET raw_payload = COALESCE(raw_payload, '{}'::jsonb) || %s::jsonb WHERE id = %s",
            (json.dumps(patch), product_id),
        )


def get_product(product_id: str) -> dict | None:
    with get_cursor() as cur:
        cur.execute("SELECT * FROM products WHERE id = %s", (product_id,))
        return cur.fetchone()


def find_active_duplicate(external_id: str) -> str | None:
    """Returns the id of an existing pending/approved row for this exact
    product (identified by its source-prefixed external_id, e.g.
    'cjdropshipping:04A22450...' or 'aliexpress:1005...'), if one exists, so
    we don't queue the same product twice. Source-prefixed so a CJ product
    and an AliExpress product can never collide on this check."""
    if not external_id:
        return None
    with get_cursor() as cur:
        cur.execute(
            """
            SELECT id FROM products
            WHERE raw_payload->>'external_id' = %s
              AND approval_status IN ('pending', 'approved')
            LIMIT 1
            """,
            (external_id,),
        )
        row = cur.fetchone()
        return str(row["id"]) if row else None


def find_active_duplicates(external_ids: list) -> set:
    """Bulk version of find_active_duplicate — one round trip instead of one
    query per scraped product. Used by each scraper to skip the expensive
    variant/shipping lookup entirely for products we already know are
    duplicates, instead of burning API calls on items that are going to
    get rejected anyway."""
    ids = [i for i in external_ids if i]
    if not ids:
        return set()
    with get_cursor() as cur:
        cur.execute(
            """
            SELECT DISTINCT raw_payload->>'external_id' AS ext_id FROM products
            WHERE raw_payload->>'external_id' = ANY(%s)
              AND approval_status IN ('pending', 'approved')
            """,
            (ids,),
        )
        return {row["ext_id"] for row in cur.fetchall()}


def list_products(approval_status: str = None, status: str = None, run_id: str = None,
                   limit: int = 200) -> list[dict]:
    query = "SELECT * FROM products WHERE 1=1"
    params = []
    if approval_status:
        query += " AND approval_status = %s"
        params.append(approval_status)
    if status:
        query += " AND status = %s"
        params.append(status)
    if run_id:
        query += " AND run_id = %s"
        params.append(run_id)
    query += " ORDER BY created_at DESC LIMIT %s"
    params.append(limit)

    with get_cursor() as cur:
        cur.execute(query, params)
        return cur.fetchall()


def latest_run_id() -> str | None:
    with get_cursor() as cur:
        cur.execute(
            "SELECT run_id FROM products WHERE run_id IS NOT NULL "
            "ORDER BY created_at DESC LIMIT 1"
        )
        row = cur.fetchone()
        return str(row["run_id"]) if row and row["run_id"] else None


def list_runs(limit: int = 20) -> list[dict]:
    """Summary row per run_id — used by the Run History page."""
    with get_cursor() as cur:
        cur.execute(
            """
            SELECT run_id,
                   MIN(run_keyword) AS keyword,
                   MIN(run_destination_country) AS destination_country,
                   MIN(created_at) AS started_at,
                   COUNT(*) FILTER (WHERE status != 'rejected_sourcing') AS sourced_count,
                   COUNT(*) FILTER (WHERE status = 'rejected_sourcing') AS rejected_count,
                   COUNT(*) FILTER (WHERE approval_status = 'approved') AS approved_count
            FROM products
            WHERE run_id IS NOT NULL
            GROUP BY run_id
            ORDER BY started_at DESC
            LIMIT %s
            """,
            (limit,),
        )
        return cur.fetchall()


def approve_product(product_id: str) -> bool:
    with get_cursor(commit=True) as cur:
        cur.execute(
            "UPDATE products SET approval_status = 'approved' WHERE id = %s RETURNING id",
            (product_id,),
        )
        return cur.fetchone() is not None

def find_product_by_request_id(request_id: str) -> dict | None:
    with get_cursor() as cur:
        cur.execute(
            """
            SELECT * FROM products
            WHERE raw_payload->>'higgsfield_request_id' = %s
            LIMIT 1
            """,
            (request_id,),
        )
        return cur.fetchone()

# ---------------------------------------------------------------------- #
# Agent audit log
# ---------------------------------------------------------------------- #

def log_agent_action(product_id: str, agent_name: str, action_taken: str, critique: str = None) -> None:
    with get_cursor(commit=True) as cur:
        cur.execute(
            "INSERT INTO agent_logs (product_id, agent_name, action_taken, critique) "
            "VALUES (%s, %s, %s, %s)",
            (product_id, agent_name, action_taken, critique),
        )


def get_logs_for_product(product_id: str) -> list[dict]:
    with get_cursor() as cur:
        cur.execute(
            "SELECT * FROM agent_logs WHERE product_id = %s ORDER BY created_at ASC",
            (product_id,),
        )
        return cur.fetchall()

def check_connection() -> tuple[bool, str]:
    """Returns (True, 'Connected') or (False, error_message)."""
    try:
        with get_cursor() as cur:
            cur.execute("SELECT 1")
        return True, "Connected"
    except Exception as e:
        return False, str(e)