from __future__ import annotations

import logging
from typing import Any, Mapping

from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession, create_async_engine
from sqlalchemy.orm import declarative_base, sessionmaker
from sqlalchemy import (
    Column,
    Integer,
    String,
    JSON,
    UniqueConstraint,
    select,
    text,
)
from sqlalchemy.dialects.postgresql import insert as postgres_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert

from .config import settings

logger = logging.getLogger(__name__)

Base = declarative_base()


class Cache(Base):
    __tablename__ = "cache"

    id = Column(Integer, primary_key=True)
    ioc = Column(String, nullable=False)
    provider = Column(String, nullable=False)
    ioc_type = Column(String, nullable=False)
    raw_response = Column(JSON, nullable=False)

    __table_args__ = (UniqueConstraint("ioc", "provider", name="uix_ioc_provider"),)


class DatasetEntry(Base):
    __tablename__ = "dataset"

    id = Column(Integer, primary_key=True)
    ioc = Column(String, nullable=False)
    ioc_type = Column(String, nullable=False)
    label = Column(String, nullable=False)


class TaskEntry(Base):
    __tablename__ = "tasks"

    id = Column(String, primary_key=True)
    status = Column(String, nullable=False)
    iocs = Column(JSON, nullable=False)
    results = Column(JSON, nullable=False, default=dict)
    verdicts = Column(JSON, nullable=False, default=dict)
    processed = Column(JSON, nullable=False, default=list)
    error = Column(String, nullable=True)
    input_text = Column(String, nullable=True)


engine = create_async_engine(settings.database_url, echo=False)
SessionLocal = sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)


async def init_db() -> None:
    logger.debug("Инициализируем схему базы данных")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        if conn.dialect.name == "postgresql":
            await _sync_postgres_sequences(conn)
    logger.info("Инициализация базы данных завершена")


async def _sync_postgres_sequences(conn: AsyncConnection) -> None:
    """Синхронизирует последовательности PostgreSQL после миграции."""

    logger.debug("Синхронизируем последовательности PostgreSQL для таблиц кеша")
    table = Cache.__tablename__
    sequence_sql = text(
        f"SELECT setval("
        f"pg_get_serial_sequence('{table}', 'id'), "
        f"COALESCE((SELECT MAX(id) FROM {table}), 1), true)"
    )
    await conn.execute(sequence_sql)
    logger.debug("Последовательности PostgreSQL синхронизированы")


async def get_cached_result(
    ioc: str,
    provider: str,
) -> tuple[str, str, str, object] | None:
    logger.debug("Проверяем кеш для IOC %s через провайдера %s", ioc, provider)
    async with SessionLocal() as session:
        stmt = select(Cache).where(Cache.ioc == ioc, Cache.provider == provider)
        res = await session.execute(stmt)
        cache = res.scalars().first()
        if cache:
            logger.debug("Найден кеш для %s через %s", ioc, provider)
            return cache.ioc, cache.ioc_type, cache.provider, cache.raw_response
    logger.debug("Кеша для %s через %s не найдено", ioc, provider)
    return None


async def cache_result(
    ioc: str,
    provider: str,
    raw_payload: object,
    ioc_type: str,
) -> None:
    logger.debug(
        "Сохраняем результат для IOC %s через провайдера %s",
        ioc,
        provider,
    )
    stored_type = ioc_type
    async with SessionLocal() as session:
        stmt = select(Cache).where(Cache.ioc == ioc, Cache.provider == provider)
        res = await session.execute(stmt)
        cache = res.scalars().first()
        if cache:
            cache.raw_response = raw_payload
            cache.ioc_type = stored_type
        else:
            cache = Cache(
                ioc=ioc,
                provider=provider,
                ioc_type=stored_type,
                raw_response=raw_payload,
            )
            session.add(cache)
        await session.commit()
    logger.debug("Результат сохранён в кеше для IOC %s через %s", ioc, provider)


async def upsert_task_record(
    task_id: str,
    *,
    status: str,
    iocs: list[str],
    results: Mapping[str, Mapping[str, Any]],
    verdicts: Mapping[str, Any],
    processed: list[str],
    error: str | None,
    input_text: str | None,
) -> None:
    logger.debug("Синхронизируем запись задачи %s в БД", task_id)
    values = dict(
        id=task_id,
        status=status,
        iocs=list(iocs),
        results=dict(results),
        verdicts=dict(verdicts),
        processed=list(processed),
        error=error,
        input_text=input_text,
    )
    async with SessionLocal() as session:
        dialect = session.bind.dialect.name if session.bind else ""
        if dialect == "postgresql":
            stmt = (
                postgres_insert(TaskEntry)
                .values(**values)
                .on_conflict_do_update(index_elements=[TaskEntry.id], set_=values)
            )
            await session.execute(stmt)
        elif dialect == "sqlite":
            stmt = (
                sqlite_insert(TaskEntry)
                .values(**values)
                .on_conflict_do_update(index_elements=[TaskEntry.id], set_=values)
            )
            await session.execute(stmt)
        else:
            existing = await session.get(TaskEntry, task_id)
            if existing:
                for field, value in values.items():
                    setattr(existing, field, value)
            else:
                session.add(TaskEntry(**values))
        await session.commit()
    logger.debug("Запись задачи %s обновлена", task_id)


async def fetch_task_record(task_id: str) -> TaskEntry | None:
    logger.debug("Читаем состояние задачи %s из БД", task_id)
    async with SessionLocal() as session:
        return await session.get(TaskEntry, task_id)
