import sqlite3
import zstandard as zstd
import tempfile
import os
from typing import Any, Optional
from pathlib import Path

__all__ = ['ZstSQLiteReader']

class ZstSQLiteReader:
    def __init__(self, zst_path: str | Path):
        self.zst_path = zst_path
        self.temp_path: Optional[str] = None
        self.conn: Optional[sqlite3.Connection] = None

    def __enter__(self):
        # 1. 임시 파일 생성 (delete=False: 윈도우 파일 잠금 문제 방지)
        self.temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.db')
        self.temp_path = self.temp_file.name

        try:
            # 2. 압축 해제 (스트리밍 방식)
            dctx = zstd.ZstdDecompressor()
            with open(self.zst_path, 'rb') as zf:
                dctx.copy_stream(zf, self.temp_file)
            self.temp_file.close()  # 쓰기 완료 후 닫기

            # 3. SQLite 연결 (URI 모드, 읽기 전용)
            self.conn = sqlite3.connect(f"file:{self.temp_path}?mode=ro", uri=True)
            self.conn.row_factory = sqlite3.Row

            # 4. [성능 최적화] 읽기 전용 DB 튜닝
            # mmap_size: 메모리 맵핑 활성화 (대용량 읽기 속도 향상)
            # journal_mode=OFF: 롤백 저널 끄기 (임시 파일이라 불필요)
            # synchronous=OFF: 디스크 동기화 대기 안 함
            self.conn.executescript(
                """
                PRAGMA mmap_size = 30000000000; 
                PRAGMA journal_mode = OFF;
                PRAGMA synchronous = OFF;
                PRAGMA temp_store = MEMORY;
            """
            )

            return self

        except Exception as e:
            self._cleanup()
            raise e

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

    def _cleanup(self):
        """리소스 정리"""
        if self.conn:
            try:
                self.conn.close()
            except:
                pass
            self.conn = None

        if self.temp_path and os.path.exists(self.temp_path):
            try:
                os.remove(self.temp_path)
            except OSError:
                # 윈도우 등에서 파일이 아직 잠겨있는 경우 무시 (OS가 나중에 정리)
                pass
            self.temp_path = None

    def query(self, sql: str, params: tuple | dict = ()) -> list[sqlite3.Row]:
        """Raw SQL 실행"""
        if not self.conn:
            raise ConnectionError("DB is not connected.")

        cursor = self.conn.cursor()
        cursor.execute(sql, params)
        return cursor.fetchall()

    def fetch(
        self,
        table: str,
        columns: str = "*",
        where: str | list[str] | dict[str, Any] = None,
        params: tuple | dict = (),
        limit: Optional[int] = None
    ) -> list[sqlite3.Row]:
        """
        조건부 데이터 조회 헬퍼
        """
        query_params = params
        where_clause = ""

        # 조건절 빌드 로직 최적화
        if isinstance(where, dict):
            # dict: {'name': 'Alice'} -> WHERE name = ?
            where_clause = " AND ".join(f"{k} = ?" for k in where)
            query_params = tuple(where.values())

        elif isinstance(where, list):
            # list: ["age > ?", "name = ?"] -> WHERE age > ? AND name = ?
            where_clause = " AND ".join(where)

        elif isinstance(where, str):
            # String: "age > 20"
            where_clause = where

        # SQL 조립 (f-string 활용)
        sql_parts = [f"SELECT {columns} FROM {table}"]

        if where_clause:
            sql_parts.append(f"WHERE {where_clause}")

        if limit is not None:
            sql_parts.append(f"LIMIT {limit}")

        full_sql = " ".join(sql_parts)

        return self.query(full_sql, query_params)