from __future__ import annotations

import os
import sqlite3
import tempfile
from pathlib import Path
from typing import Any, Iterable, Optional

import zstandard as zstd

__all__ = ["ZstSQLiteReader", "compress_zst", "decompress_zst"]


def compress_zst(
    src: str | Path,
    dst: Optional[str | Path] = None,
    level: int = 3,
    chunk_size: int = 4 * 1024 * 1024,
    compressor: Optional[zstd.ZstdCompressor] = None,
    threads: Optional[int] = None,
) -> Path:
    """
    단일 파일을 Zstandard(.zst) 형식으로 압축합니다.

    :param src: 원본 파일 경로.
    :type src: str 또는 pathlib.Path
    :param dst: 출력 .zst 파일 경로.
        생략하면 ``<src>.zst`` 경로로 저장합니다.
    :type dst: str 또는 pathlib.Path 또는 None
    :param level: 압축 레벨(1–22). 값이 높을수록 압축률은 좋아지지만 속도는 느려집니다.
        ``compressor`` 인스턴스를 직접 전달하면 이 값은 무시됩니다.
    :type level: int
    :param chunk_size: 파일을 읽고 쓰는 단위 크기(바이트).
    :type chunk_size: int
    :param compressor: 재사용을 위해 미리 생성한 :class:`zstandard.ZstdCompressor` 인스턴스.
        여러 파일을 연속으로 압축할 때 유용합니다. 이 값을 지정하면 ``threads``,
        ``level`` 인자는 무시됩니다.
    :type compressor: zstandard.ZstdCompressor 또는 None
    :param threads: 압축에 사용할 스레드 수.
        ``None``이면 현재 시스템의 CPU 코어 수에서 1개를 뺀 값을 사용합니다.
        (예: 8코어 시스템이면 기본값은 7 스레드, 최소 1 이상)
        1을 지정하면 단일 스레드로 동작합니다.
    :type threads: int 또는 None
    :returns: 생성된 .zst 파일의 경로.
    :rtype: pathlib.Path
    :raises FileNotFoundError: 원본 파일이 존재하지 않을 때.
    :raises ValueError: 원본 경로가 파일이 아닐 때.
    """
    src_path = Path(src)

    if not src_path.is_file():
        if not src_path.exists():
            raise FileNotFoundError(f"Source file not found: {src_path}")
        raise ValueError(f"Source path is not a file: {src_path}")

    dst_path = Path(dst) if dst is not None else src_path.with_name(src_path.name + ".zst")

    if compressor is not None:
        cctx = compressor
    else:
        if threads is None:
            cpu_count = os.cpu_count() or 1
            # 한 코어는 남겨 두고, 최소 1 스레드는 사용
            threads = max(cpu_count - 1, 1)
        cctx = zstd.ZstdCompressor(level=level, threads=threads)

    with src_path.open("rb") as f_src, dst_path.open("wb") as f_dst:
        with cctx.stream_writer(f_dst) as writer:
            for chunk in iter(lambda: f_src.read(chunk_size), b""):
                writer.write(chunk)

    return dst_path


def decompress_zst(
    src: str | Path,
    dst: Optional[str | Path] = None,
    chunk_size: int = 4 * 1024 * 1024,
    decompressor: Optional[zstd.ZstdDecompressor] = None,
) -> Path:
    """
    Zstandard(.zst) 파일을 원본 형태로 압축 해제합니다.

    :param src: .zst 압축 파일 경로.
    :type src: str 또는 pathlib.Path
    :param dst: 압축 해제 결과를 저장할 파일 경로.
        생략하면 ``<src>``에서 끝의 ``.zst`` 확장자를 제거한 이름을 사용합니다.
        확장자가 ``.zst``가 아니면 ``<src>.out`` 형태로 저장합니다.
    :type dst: str 또는 pathlib.Path 또는 None
    :param chunk_size: 파일을 읽고 쓰는 단위 크기(바이트).
    :type chunk_size: int
    :param decompressor: 재사용을 위해 미리 생성한 :class:`zstandard.ZstdDecompressor`
        인스턴스. 여러 파일을 연속으로 해제할 때 유용합니다.
    :type decompressor: zstandard.ZstdDecompressor 또는 None
    :returns: 압축 해제된 파일의 경로.
    :rtype: pathlib.Path
    :raises FileNotFoundError: 입력 파일이 존재하지 않을 때.
    :raises ValueError: 입력 경로가 파일이 아닐 때.
    """
    src_path = Path(src)

    if not src_path.is_file():
        if not src_path.exists():
            raise FileNotFoundError(f"Source file not found: {src_path}")
        raise ValueError(f"Source path is not a file: {src_path}")

    if dst is None:
        name = src_path.name
        dst_name = name[:-4] if name.endswith(".zst") else name + ".out"
        dst_path = src_path.with_name(dst_name)
    else:
        dst_path = Path(dst)

    dctx = decompressor or zstd.ZstdDecompressor()

    with src_path.open("rb") as f_src, dst_path.open("wb") as f_dst:
        with dctx.stream_reader(f_src) as reader:
            for chunk in iter(lambda: reader.read(chunk_size), b""):
                f_dst.write(chunk)

    return dst_path


class ZstSQLiteReader:
    """
    Zstandard로 압축된 SQLite 데이터베이스 파일을 읽기 위한 컨텍스트 매니저입니다.

    압축된 ``.zst`` 파일을 임시 파일로 압축 해제한 뒤, 해당 임시 파일을
    SQLite 데이터베이스로 엽니다. 컨텍스트가 종료되면 임시 파일은 자동으로
    삭제됩니다.
    """

    __slots__ = ("zst_path", "temp_path", "conn")

    def __init__(self, zst_path: str | Path):
        """
        :param zst_path: Zstandard로 압축된 SQLite 데이터베이스 파일 경로.
        :type zst_path: str 또는 pathlib.Path
        """
        self.zst_path = str(zst_path)
        self.temp_path: Optional[str] = None
        self.conn: Optional[sqlite3.Connection] = None

    def __enter__(self) -> "ZstSQLiteReader":
        self._open()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self._cleanup()

    def _open(self) -> None:
        """
        내부용: 압축된 DB를 임시 파일로 해제하고 SQLite 연결을 생성합니다.
        """
        # 임시 파일 생성
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
        self.temp_path = tmp.name

        try:
            # 압축 해제
            dctx = zstd.ZstdDecompressor()
            with open(self.zst_path, "rb") as zf:
                dctx.copy_stream(zf, tmp)

            # flush 및 fsync로 디스크 반영 보장 (특히 Windows에서)
            tmp.flush()
            os.fsync(tmp.fileno())
            tmp.close()

            # SQLite 연결
            self.conn = sqlite3.connect(self.temp_path)
            self.conn.row_factory = sqlite3.Row

            # 임시 DB이므로 성능 위주 설정
            self.conn.executescript(
                """
                PRAGMA journal_mode = OFF;
                PRAGMA synchronous = OFF;
                PRAGMA temp_store = MEMORY;
                PRAGMA cache_size = -20000;
                """
            )
        except Exception:
            # 예외 발생 시 리소스 정리 후 예외 재전파
            self._cleanup()
            raise

    def _cleanup(self) -> None:
        """
        내부용: SQLite 연결과 임시 파일을 정리합니다.
        """
        if self.conn is not None:
            try:
                self.conn.close()
            except Exception:
                pass
            self.conn = None

        if self.temp_path is not None:
            try:
                if os.path.exists(self.temp_path):
                    os.remove(self.temp_path)
            except OSError:
                # Windows에서 파일이 잠겨 있을 수 있으므로 삭제 실패는 무시합니다.
                pass
            self.temp_path = None

    # --- 쿼리 헬퍼 -----------------------------------------------------

    def _ensure_conn(self) -> sqlite3.Connection:
        """
        내부용: 연결이 유효한지 검사하고, 없으면 예외를 발생시킵니다.
        """
        if self.conn is None:
            raise ConnectionError("SQLite 데이터베이스가 연결되어 있지 않습니다.")
        return self.conn

    def query(
        self,
        sql: str,
        params: Iterable[Any] | dict[str, Any] = (),
    ) -> list[sqlite3.Row]:
        """
        SQL 쿼리를 실행하고 모든 결과 행을 반환합니다.

        :param sql: 실행할 SQL 문자열.
        :type sql: str
        :param params: 쿼리에 바인딩할 파라미터.
            위치 기반(iterable) 또는 이름 기반(dict) 모두 사용할 수 있습니다.
        :type params: iterable 또는 dict
        :returns: 조회된 결과 행들의 리스트.
        :rtype: list[sqlite3.Row]
        """
        conn = self._ensure_conn()
        cur = conn.cursor()
        cur.execute(sql, params)
        rows = cur.fetchall()
        cur.close()
        return rows

    def fetch(
        self,
        table: str,
        columns: str | Iterable[str] = "*",
        where: Optional[str | Iterable[str] | dict[str, Any]] = None,
        params: Iterable[Any] | dict[str, Any] = (),
        limit: Optional[int] = None,
    ) -> list[sqlite3.Row]:
        """
        간단한 ``SELECT`` 조회를 위한 헬퍼 메서드입니다.

        :param table: 조회할 테이블 이름.
        :type table: str
        :param columns: 조회할 컬럼 목록.
            ``"*"`` 또는 컬럼 이름들의 iterable을 사용할 수 있습니다.
        :type columns: str 또는 iterable
        :param where: WHERE 조건을 지정합니다. 다음 형태를 지원합니다.

            * ``str``:
              SQL WHERE 절 문자열을 그대로 사용합니다.
              예: ``"a = 1 OR b = 2"`` 처럼 ``OR`` 또는 괄호 등을 포함한
              복잡한 조건을 직접 작성할 수 있습니다.

            * 문자열의 iterable (예: list, tuple, set):
              각 문자열을 하나의 조건으로 간주하고, 이들을 ``AND`` 로 연결합니다.
              예: ``["a = 1", "b = 2"]`` → ``"a = 1 AND b = 2"``

              이때 각 조건 문자열 내부에는 ``OR`` 등을 자유롭게 포함할 수 있습니다.
              예: ``["a = 1 OR b = 2", "c = 3"]``
              → ``"a = 1 OR b = 2 AND c = 3"``

            * ``dict``:
              키를 컬럼명으로 간주하고, 각 항목을 ``"column = ?"`` 형태로 만들며,
              값들은 위치 파라미터로 전달됩니다. 항목들은 ``AND`` 로 연결됩니다.
              예: ``{"a": 1, "b": 2}`` → ``"a = ? AND b = ?"``

              이 방식은 SQL 인젝션 방지에 유리하지만,
              연산자는 ``=`` 로 고정되며, ``OR`` 같은 복잡한 조건은 직접 표현할 수 없습니다.

        :type where: str 또는 iterable 또는 dict 또는 None
        :param params: 쿼리 파라미터.
            ``where``를 문자열로 쓸 때는 여기에 대응하는 파라미터를 넘겨야 합니다.
            ``where``가 dict일 때는 무시되며, dict의 값이 자동으로 사용됩니다.
        :type params: iterable 또는 dict
        :param limit: SELECT 결과의 최대 행 수(LIMIT).
        :type limit: int 또는 None
        :returns: 조회된 결과 행들의 리스트.
        :rtype: list[sqlite3.Row]
        """
        # 컬럼 문자열 구성
        if isinstance(columns, str):
            col_str = columns
        else:
            col_str = ", ".join(columns)

        where_clause = ""
        query_params: Iterable[Any] | dict[str, Any] = params

        if isinstance(where, dict):
            where_clause = " AND ".join(f"{k} = ?" for k in where)
            query_params = tuple(where.values())
        elif isinstance(where, (list, tuple, set)):
            where_list = list(where)
            if where_list:
                where_clause = " AND ".join(where_list)
        elif isinstance(where, str):
            where_clause = where or ""

        sql_parts = [f"SELECT {col_str} FROM {table}"]
        if where_clause:
            sql_parts.append(f"WHERE {where_clause}")
        if limit is not None:
            sql_parts.append(f"LIMIT {limit}")

        sql = " ".join(sql_parts)
        return self.query(sql, query_params)