"""
PostgreSQL sink for extractor.

Streams records into a Postgres table using COPY or executemany.
JSON columns for nested fields (entities, relations, classification).

Auth (env vars):
    EXTRACTOR_PG_DSN    (full DSN: postgresql://user:pass@host:5432/db)
    or standard libpq env vars: PGHOST, PGPORT, PGUSER, PGPASSWORD, PGDATABASE

Table DDL (auto-created):
    CREATE TABLE IF NOT EXISTS {table} (
        id TEXT PRIMARY KEY,
        element_type TEXT,
        text TEXT,
        source_filename TEXT,
        source_sha256 TEXT,
        source_uri TEXT,
        source_media_type TEXT,
        page INTEGER,
        sequence_index INTEGER,
        section_path TEXT[],
        schema_version TEXT,
        extraction_method TEXT,
        confidence REAL,
        token_count INTEGER,
        entities JSONB DEFAULT '[]',
        classification JSONB,
        relations JSONB DEFAULT '[]',
        structured_fields JSONB,
        table_data JSONB,
        created_at TIMESTAMPTZ DEFAULT now()
    )

Install: pip install 'coresdk-extractor[postgres]'
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
import re
from typing import Any

from .base import BaseSink
from .registry import register_sink

_TABLE_NAME_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]{0,62}$")

logger = logging.getLogger(__name__)


class PostgresSink(BaseSink):
    def __init__(
        self,
        dsn: str | None = None,
        table: str = "extractor_elements",
        batch_size: int = 500,
        upsert: bool = True,
    ) -> None:
        self.dsn = dsn or os.environ.get("EXTRACTOR_PG_DSN", "")
        if not _TABLE_NAME_RE.match(table):
            raise ValueError(
                f"Invalid table name {table!r}. Must match ^[a-zA-Z_][a-zA-Z0-9_]{{0,62}}$"
            )
        self.table = table
        self.batch_size = batch_size
        self.upsert = upsert
        self._conn = None
        self._buffer: list[dict[str, Any]] = []
        self._total_inserted = 0

    def connect(self) -> None:
        try:
            import psycopg2
            import psycopg2.extras
        except ImportError as e:
            raise ImportError(
                "psycopg2 not installed. Run: pip install 'coresdk-extractor[postgres]'"
            ) from e
        self._conn = psycopg2.connect(self.dsn) if self.dsn else psycopg2.connect()
        self._conn.autocommit = False
        with self._conn.cursor() as cur:
            cur.execute(f"""
                CREATE TABLE IF NOT EXISTS {self.table} (
                    id TEXT PRIMARY KEY,
                    element_type TEXT,
                    text TEXT,
                    source_filename TEXT,
                    source_sha256 TEXT,
                    source_uri TEXT,
                    source_media_type TEXT,
                    page INTEGER,
                    sequence_index INTEGER,
                    section_path TEXT[],
                    schema_version TEXT,
                    extraction_method TEXT,
                    confidence REAL,
                    token_count INTEGER,
                    entities JSONB DEFAULT '[]',
                    classification JSONB,
                    relations JSONB DEFAULT '[]',
                    structured_fields JSONB,
                    table_data JSONB,
                    created_at TIMESTAMPTZ DEFAULT now()
                )
            """)
            cur.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table}_sha256 ON {self.table}(source_sha256)"
            )
            cur.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table}_type ON {self.table}(element_type)"
            )
        self._conn.commit()
        logger.info("PostgresSink connected: table=%s", self.table)

    def feed(self, record: dict[str, Any]) -> None:
        self._buffer.append(record)
        if len(self._buffer) >= self.batch_size:
            self.write_batch(self._buffer)
            self._buffer = []

    def write_batch(self, records: list[dict[str, Any]]) -> None:
        if not records or self._conn is None:
            return
        import psycopg2.extras

        rows = [self._flatten(r) for r in records]
        cols = list(rows[0].keys())
        if self.upsert:
            set_clause = ", ".join(f"{c} = EXCLUDED.{c}" for c in cols if c != "id")
            sql = (
                f"INSERT INTO {self.table} ({', '.join(cols)}) VALUES %s "
                f"ON CONFLICT (id) DO UPDATE SET {set_clause}"
            )
        else:
            sql = (
                f"INSERT INTO {self.table} ({', '.join(cols)}) VALUES %s "
                f"ON CONFLICT (id) DO NOTHING"
            )
        values = [[r[c] for c in cols] for r in rows]
        with self._conn.cursor() as cur:
            psycopg2.extras.execute_values(cur, sql, values, page_size=self.batch_size)
        self._conn.commit()
        self._total_inserted += len(rows)

    def _flatten(self, record: dict[str, Any]) -> dict[str, Any]:
        def _j(v: Any) -> str | None:
            return json.dumps(v, ensure_ascii=False) if v is not None else None

        return {
            "id": record.get("id", ""),
            "element_type": record.get("element_type", ""),
            "text": record.get("text", ""),
            "source_filename": record.get("source_filename", ""),
            "source_sha256": record.get("source_sha256", ""),
            "source_uri": record.get("source_uri"),
            "source_media_type": record.get("source_media_type", ""),
            "page": record.get("page"),
            "sequence_index": record.get("sequence_index", 0),
            "section_path": record.get("section_path", []),
            "schema_version": record.get("schema_version", ""),
            "extraction_method": record.get("extraction_method", ""),
            "confidence": record.get("confidence"),
            "token_count": record.get("token_count"),
            "entities": _j(record.get("entities")),
            "classification": _j(record.get("classification")),
            "relations": _j(record.get("relations")),
            "structured_fields": _j(record.get("structured_fields")),
            "table_data": _j(record.get("table")),
        }

    def flush(self) -> None:
        if self._buffer:
            self.write_batch(self._buffer)
            self._buffer = []

    def close(self) -> None:
        if self._conn:
            with contextlib.suppress(Exception):
                self._conn.close()
            self._conn = None

    @property
    def total_inserted(self) -> int:
        return self._total_inserted


register_sink("postgres", PostgresSink)
register_sink("postgresql", PostgresSink)
