from loguru import logger
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

from session_hub_core.config import DatabaseConfig
from session_hub_core.models import Base

_SESSIONS_TSV_FUNCTION_SQL = """
CREATE OR REPLACE FUNCTION sessions_content_tsv_update()
RETURNS trigger AS $$
BEGIN
  NEW.content_tsv :=
    setweight(to_tsvector('english', unaccent(coalesce(NEW.title, ''))), 'A') ||
    setweight(to_tsvector('english', unaccent(coalesce(
      (SELECT string_agg(content, ' ')
       FROM messages WHERE messages.session_pk = NEW.id),
      ''
    ))), 'B');
  RETURN NEW;
END;
$$ LANGUAGE plpgsql;
"""

_SESSIONS_TSV_TRIGGER_SQL = (
    "CREATE TRIGGER sessions_content_tsv_trigger\n"
    "  BEFORE INSERT OR UPDATE OF title ON sessions\n"
    "  FOR EACH ROW EXECUTE FUNCTION sessions_content_tsv_update();\n"
)
_SESSIONS_TSV_TRIGGER_DROP_SQL = "DROP TRIGGER IF EXISTS sessions_content_tsv_trigger ON sessions;"

_MESSAGES_SESSION_TSV_FUNCTION_SQL = """
CREATE OR REPLACE FUNCTION messages_session_tsv_update()
RETURNS trigger AS $$
DECLARE
  v_session_pk INTEGER;
BEGIN
  v_session_pk := COALESCE(NEW.session_pk, OLD.session_pk);
  UPDATE sessions SET content_tsv =
    setweight(to_tsvector('english', unaccent(coalesce(title, ''))), 'A') ||
    setweight(to_tsvector('english', unaccent(coalesce(
      (SELECT string_agg(content, ' ') FROM messages WHERE session_pk = v_session_pk), ''
    ))), 'B')
  WHERE id = v_session_pk;
  RETURN COALESCE(NEW, OLD);
END;
$$ LANGUAGE plpgsql;
"""

_MESSAGES_SESSION_TSV_TRIGGER_SQL = (
    "CREATE TRIGGER messages_session_tsv_trigger\n"
    "  AFTER INSERT OR UPDATE OR DELETE ON messages\n"
    "  FOR EACH ROW EXECUTE FUNCTION messages_session_tsv_update();\n"
)
_MESSAGES_SESSION_TSV_TRIGGER_DROP_SQL = (
    "DROP TRIGGER IF EXISTS messages_session_tsv_trigger ON messages;"
)

_MESSAGES_CONTENT_TSV_FUNCTION_SQL = """
CREATE OR REPLACE FUNCTION messages_content_tsv_update()
RETURNS trigger AS $$
BEGIN
  NEW.content_tsv := to_tsvector('english', unaccent(coalesce(NEW.content, '')));
  RETURN NEW;
END;
$$ LANGUAGE plpgsql;
"""

_MESSAGES_CONTENT_TSV_TRIGGER_SQL = (
    "CREATE TRIGGER messages_content_tsv_trigger\n"
    "  BEFORE INSERT OR UPDATE OF content ON messages\n"
    "  FOR EACH ROW EXECUTE FUNCTION messages_content_tsv_update();\n"
)
_MESSAGES_CONTENT_TSV_TRIGGER_DROP_SQL = (
    "DROP TRIGGER IF EXISTS messages_content_tsv_trigger ON messages;"
)

_PARTMAN_CREATE_PARENT_SQL = """
SELECT partman.create_parent(
  p_parent_table => 'public.sync_log',
  p_control => 'started_at',
  p_type => 'native',
  p_interval => 'monthly',
  p_premake => 4
)
"""

_PARTMAN_RETENTION_SQL = """
UPDATE partman.part_config
SET retention = '12 months', retention_keep_table = false
WHERE parent_table = 'public.sync_log'
"""

_PG_CRON_CLEANUP_SQL = """
SELECT cron.schedule(
  'cleanup-stale-sync-log',
  '0 * * * *',
  $$UPDATE sync_log SET status = 'failed', error = 'interrupted'
    WHERE status = 'running' AND started_at < now() - interval '1 hour'$$
)
"""

_PG_CRON_PARTMAN_SQL = """
SELECT cron.schedule(
  'partman-maintenance',
  '0 3 * * *',
  $$CALL partman.run_maintenance_proc()$$
)
"""


async def create_schema(config: DatabaseConfig) -> None:
    engine = create_async_engine(
        config.dsn, pool_size=config.pool_size, max_overflow=config.max_overflow, echo=config.echo
    )
    try:
        async with engine.begin() as conn:
            await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
            await conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
            await conn.execute(text("CREATE EXTENSION IF NOT EXISTS unaccent"))
            await conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_stat_statements"))
            available = {
                row[0]
                for row in (
                    await conn.execute(
                        text(
                            "SELECT name FROM pg_available_extensions "
                            "WHERE name IN ('pg_cron', 'pg_partman', 'pg_repack')"
                        )
                    )
                ).all()
            }
            for ext in ("pg_cron", "pg_partman", "pg_repack"):
                if ext not in available:
                    logger.warning("Extension '{}' not available; skipping", ext)
                    continue
                try:
                    await conn.execute(text(f"CREATE EXTENSION IF NOT EXISTS {ext}"))
                except Exception as exc:
                    logger.warning("Could not create {} extension: {}", ext, exc)
            await conn.run_sync(Base.metadata.create_all)
            await conn.execute(text(_SESSIONS_TSV_FUNCTION_SQL))
            await conn.execute(text(_SESSIONS_TSV_TRIGGER_DROP_SQL))
            await conn.execute(text(_SESSIONS_TSV_TRIGGER_SQL))
            await conn.execute(text(_MESSAGES_SESSION_TSV_FUNCTION_SQL))
            await conn.execute(text(_MESSAGES_SESSION_TSV_TRIGGER_DROP_SQL))
            await conn.execute(text(_MESSAGES_SESSION_TSV_TRIGGER_SQL))
            await conn.execute(text(_MESSAGES_CONTENT_TSV_FUNCTION_SQL))
            await conn.execute(text(_MESSAGES_CONTENT_TSV_TRIGGER_DROP_SQL))
            await conn.execute(text(_MESSAGES_CONTENT_TSV_TRIGGER_SQL))
            if "pg_partman" in available:
                try:
                    await conn.execute(text(_PARTMAN_CREATE_PARENT_SQL))
                    await conn.execute(text(_PARTMAN_RETENTION_SQL))
                except Exception as exc:
                    logger.warning("pg_partman partitioning setup failed: {}", exc)
            if "pg_cron" in available:
                try:
                    await conn.execute(text(_PG_CRON_CLEANUP_SQL))
                    await conn.execute(text(_PG_CRON_PARTMAN_SQL))
                except Exception as exc:
                    logger.warning("pg_cron job scheduling failed: {}", exc)
    finally:
        await engine.dispose()
