"""
database.py
===========
SQLite mirror of the metadata CSV, used only for fast searching/filtering in
the Dataset Explorer tab. The CSV (see ``metadata.py``) remains the master
export — this database is a derived, disposable convenience index.

Why "download, modify, re-upload" instead of querying ``hf://`` directly?
--------------------------------------------------------------------------
``sqlite3`` needs a real local file handle with support for locking/seeking,
which object storage / fsspec paths don't reliably provide. So each mutating
operation:

  1. Pulls the current DB file from the bucket to a local scratch path
     (skipped if using a mounted volume, where the "local" path *is* already
     inside the bucket, i.e. zero-copy).
  2. Opens it with ``sqlite3``, applies the change.
  3. Immediately pushes the updated file back to the bucket.

All of this is wrapped in a process-wide lock so concurrent uploads cannot
interleave and corrupt the database file.
"""

from __future__ import annotations

import logging
import os
import sqlite3
import threading
from contextlib import contextmanager
from typing import Dict, Iterator, List

from bucket import BucketStorage
from utils import DATABASE_PATH, METADATA_COLUMNS

logger = logging.getLogger("speech_collector.database")

_db_lock = threading.RLock()

_CREATE_TABLE_SQL = f"""
CREATE TABLE IF NOT EXISTS samples (
    {", ".join(f'{col} TEXT' for col in METADATA_COLUMNS if col not in ("sample_id", "age", "audio_duration", "sample_rate"))},
    sample_id TEXT PRIMARY KEY,
    age INTEGER,
    audio_duration REAL,
    sample_rate INTEGER
)
"""


class SpeechDatabase:
    """Thread-safe SQLite mirror living inside the bucket."""

    def __init__(self, bucket: BucketStorage, local_scratch_path: str = "/tmp/speech_dataset_local.db") -> None:
        self.bucket = bucket
        self.local_scratch_path = local_scratch_path
        self._ensure_db_exists()

    def _ensure_db_exists(self) -> None:
        with _db_lock:
            if self.bucket.exists(DATABASE_PATH):
                self.bucket.download_to_local(DATABASE_PATH, self.local_scratch_path)
            else:
                logger.info("speech_dataset.db not found in bucket — creating a new one.")
                if os.path.exists(self.local_scratch_path):
                    os.remove(self.local_scratch_path)
            conn = sqlite3.connect(self.local_scratch_path)
            try:
                conn.execute(_CREATE_TABLE_SQL)
                conn.commit()
            finally:
                conn.close()
            self.bucket.sync_local_to_bucket(self.local_scratch_path, DATABASE_PATH)

    @contextmanager
    def _connect_synced(self) -> Iterator[sqlite3.Connection]:
        """Pull latest DB from bucket, yield a connection, push changes back."""
        with _db_lock:
            if self.bucket.exists(DATABASE_PATH):
                self.bucket.download_to_local(DATABASE_PATH, self.local_scratch_path)
            conn = sqlite3.connect(self.local_scratch_path)
            try:
                yield conn
                conn.commit()
            finally:
                conn.close()
            self.bucket.sync_local_to_bucket(self.local_scratch_path, DATABASE_PATH)

    def insert_sample(self, row: Dict[str, str]) -> None:
        columns = METADATA_COLUMNS
        placeholders = ", ".join("?" for _ in columns)
        values = [row.get(col, "") for col in columns]
        sql = f"INSERT OR REPLACE INTO samples ({', '.join(columns)}) VALUES ({placeholders})"
        try:
            with self._connect_synced() as conn:
                conn.execute(sql, values)
            logger.info("Mirrored sample %s into SQLite database.", row.get("sample_id"))
        except Exception as exc:
            # The CSV row is already safely written by this point; a DB
            # mirroring failure must not be treated as a fatal upload error.
            logger.error("Failed to mirror sample into SQLite (non-fatal): %s", exc)

    def fetch_all(self) -> List[tuple]:
        try:
            with self._connect_synced() as conn:
                cursor = conn.execute(
                    "SELECT sample_id, name, gender, age, region, state, country, "
                    "language, dialect, prompt_text, transcript, audio_file, "
                    "audio_duration, timestamp FROM samples ORDER BY sample_id DESC"
                )
                return cursor.fetchall()
        except Exception as exc:
            logger.error("Failed to fetch samples from SQLite: %s", exc)
            return []

    def aggregate_stats(self) -> Dict:
        try:
            with self._connect_synced() as conn:
                total_samples, total_speakers, total_duration = conn.execute(
                    "SELECT COUNT(*), COUNT(DISTINCT name), COALESCE(SUM(audio_duration), 0) FROM samples"
                ).fetchone()

                def group_counts(column: str) -> List[tuple]:
                    return conn.execute(
                        f"SELECT {column}, COUNT(*) FROM samples "
                        f"WHERE {column} IS NOT NULL AND {column} != '' "
                        f"GROUP BY {column} ORDER BY COUNT(*) DESC"
                    ).fetchall()

                latest_uploads = conn.execute(
                    "SELECT sample_id, name, language, timestamp FROM samples "
                    "ORDER BY sample_id DESC LIMIT 10"
                ).fetchall()

                age_stats = conn.execute(
                    "SELECT MIN(age), MAX(age), AVG(age) FROM samples WHERE age IS NOT NULL"
                ).fetchone()

                return {
                    "total_samples": total_samples or 0,
                    "total_speakers": total_speakers or 0,
                    "total_duration_seconds": total_duration or 0.0,
                    "gender_distribution": group_counts("gender"),
                    "language_distribution": group_counts("language"),
                    "dialect_distribution": group_counts("dialect"),
                    "latest_uploads": latest_uploads,
                    "age_min": age_stats[0] if age_stats else None,
                    "age_max": age_stats[1] if age_stats else None,
                    "age_avg": round(age_stats[2], 1) if age_stats and age_stats[2] else None,
                }
        except Exception as exc:
            logger.error("Failed to aggregate stats from SQLite: %s", exc)
            return {
                "total_samples": 0, "total_speakers": 0, "total_duration_seconds": 0.0,
                "gender_distribution": [], "language_distribution": [], "dialect_distribution": [],
                "latest_uploads": [], "age_min": None, "age_max": None, "age_avg": None,
            }