"""Shared Postgres translation storage."""

import hashlib
import logging
from datetime import datetime, timezone

logger = logging.getLogger(__name__)


TABLE_NAME = "transduck_translations"


def _content_hash(source_text: str, project_context_hash: str, string_context_hash: str) -> str:
    return hashlib.sha256(
        (source_text + "\x00" + project_context_hash + "\x00" + string_context_hash).encode()
    ).hexdigest()


_CREATE_TABLE = f"""
CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
    source_lang TEXT NOT NULL,
    target_lang TEXT NOT NULL,
    content_hash TEXT NOT NULL,
    plural_category TEXT NOT NULL DEFAULT '',
    source_text TEXT NOT NULL,
    translated_text TEXT NOT NULL,
    model TEXT NOT NULL,
    status TEXT NOT NULL,
    created_at TEXT NOT NULL,
    project_context_hash TEXT NOT NULL,
    string_context_hash TEXT NOT NULL,
    string_context TEXT NOT NULL DEFAULT '',
    PRIMARY KEY (source_lang, target_lang, content_hash, plural_category)
)
"""


class SharedStore:
    """Postgres-backed translation store with the same interface as TranslationStore.

    Automatically reconnects when the connection is dropped by the server
    (common with managed Postgres providers that kill idle connections).
    """

    def __init__(self, url: str):
        try:
            import psycopg2
        except ImportError:
            raise ImportError(
                "psycopg2 is required for shared Postgres storage. "
                "Install it with: pip install transduck[postgres]"
            )
        self._psycopg2 = psycopg2
        self._url = url
        self._conn = None

    def _get_conn(self):
        """Return a live connection, reconnecting if the previous one was dropped."""
        if self._conn is None or self._conn.closed:
            logger.debug("Connecting to shared Postgres store")
            self._conn = self._psycopg2.connect(self._url)
            self._conn.autocommit = False
        return self._conn

    def _reconnect(self):
        """Force a fresh connection (used after a stale connection error)."""
        try:
            if self._conn is not None and not self._conn.closed:
                self._conn.close()
        except Exception:
            pass
        self._conn = None
        return self._get_conn()

    def _execute_with_retry(self, fn):
        """Execute fn(conn) with one automatic reconnect on connection errors."""
        try:
            return fn(self._get_conn())
        except (self._psycopg2.InterfaceError, self._psycopg2.OperationalError) as e:
            logger.warning("Shared store connection lost, reconnecting: %s", e)
            return fn(self._reconnect())

    def initialize(self):
        conn = self._get_conn()
        with conn.cursor() as cur:
            cur.execute(_CREATE_TABLE)
        conn.commit()

    def lookup(
        self,
        source_text: str,
        source_lang: str,
        target_lang: str,
        project_context_hash: str,
        string_context_hash: str,
    ) -> str | None:
        chash = _content_hash(source_text, project_context_hash, string_context_hash)

        def _do(conn):
            with conn.cursor() as cur:
                cur.execute(
                    f"SELECT translated_text FROM {TABLE_NAME} "
                    "WHERE source_lang=%s AND target_lang=%s AND content_hash=%s "
                    "AND plural_category='' AND status='translated'",
                    (source_lang, target_lang, chash),
                )
                row = cur.fetchone()
            if row is None:
                with conn.cursor() as cur:
                    cur.execute(
                        f"SELECT 1 FROM {TABLE_NAME} "
                        "WHERE source_text=%s AND source_lang=%s AND target_lang=%s "
                        "AND status='translated' LIMIT 1",
                        (source_text, source_lang, target_lang),
                    )
                    if cur.fetchone() is not None:
                        logger.debug(
                            "Translation exists in shared store for '%s' (%s→%s) but with a "
                            "different project context hash. Check project.context in transduck.yaml.",
                            source_text, source_lang, target_lang,
                        )
                return None
            return row[0]

        return self._execute_with_retry(_do)

    def insert(
        self,
        source_text: str,
        source_lang: str,
        target_lang: str,
        project_context_hash: str,
        string_context_hash: str,
        translated_text: str,
        model: str,
        status: str = "translated",
        string_context: str = "",
    ):
        chash = _content_hash(source_text, project_context_hash, string_context_hash)

        def _do(conn):
            with conn.cursor() as cur:
                cur.execute(
                    f"INSERT INTO {TABLE_NAME} "
                    "(source_lang, target_lang, content_hash, plural_category, source_text, "
                    "translated_text, model, status, created_at, project_context_hash, "
                    "string_context_hash, string_context) "
                    "VALUES (%s, %s, %s, '', %s, %s, %s, %s, %s, %s, %s, %s) "
                    "ON CONFLICT DO NOTHING",
                    (source_lang, target_lang, chash, source_text, translated_text, model,
                     status, datetime.now(timezone.utc).isoformat(), project_context_hash,
                     string_context_hash, string_context or ""),
                )
            conn.commit()

        self._execute_with_retry(_do)

    def lookup_plural(
        self,
        source_text: str,
        source_lang: str,
        target_lang: str,
        project_context_hash: str,
        string_context_hash: str,
    ) -> dict[str, str]:
        chash = _content_hash(source_text, project_context_hash, string_context_hash)

        def _do(conn):
            with conn.cursor() as cur:
                cur.execute(
                    f"SELECT plural_category, translated_text FROM {TABLE_NAME} "
                    "WHERE source_lang=%s AND target_lang=%s AND content_hash=%s "
                    "AND plural_category!='' AND status='translated'",
                    (source_lang, target_lang, chash),
                )
                rows = cur.fetchall()
            return {row[0]: row[1] for row in rows}

        return self._execute_with_retry(_do)

    def insert_plural(
        self,
        source_text: str,
        source_lang: str,
        target_lang: str,
        project_context_hash: str,
        string_context_hash: str,
        plural_category: str,
        translated_text: str,
        model: str,
        status: str = "translated",
        string_context: str = "",
    ):
        chash = _content_hash(source_text, project_context_hash, string_context_hash)

        def _do(conn):
            with conn.cursor() as cur:
                cur.execute(
                    f"INSERT INTO {TABLE_NAME} "
                    "(source_lang, target_lang, content_hash, plural_category, source_text, "
                    "translated_text, model, status, created_at, project_context_hash, "
                    "string_context_hash, string_context) "
                    "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) "
                    "ON CONFLICT DO NOTHING",
                    (source_lang, target_lang, chash, plural_category, source_text,
                     translated_text, model, status, datetime.now(timezone.utc).isoformat(),
                     project_context_hash, string_context_hash, string_context or ""),
                )
            conn.commit()

        self._execute_with_retry(_do)

    def close(self):
        if self._conn is not None:
            try:
                self._conn.close()
            except Exception:
                pass
            self._conn = None
