"""Database loaders for activations and pair texts from Supabase with local caching."""
from typing import Dict, Optional, List, Tuple
import os
import struct
import json
import torch
from .cache import get_cache_path, save_pair_texts_cache, save_activations_cache


def _get_db_connection(database_url: Optional[str] = None):
    """Get database connection with proper URL handling."""
    try:
        import psycopg2
    except ImportError:
        raise ImportError("psycopg2 required for database loading. Install with: pip install psycopg2-binary")

    db_url = database_url or os.environ.get("DATABASE_URL")
    if not db_url:
        raise ValueError("No database URL provided. Set DATABASE_URL env var or pass database_url parameter.")

    # Add sslmode if not present
    if "sslmode=" not in db_url:
        db_url += "?sslmode=require" if "?" not in db_url else "&sslmode=require"

    print(f"  DEBUG: Connecting to database...", flush=True)
    conn = psycopg2.connect(db_url)
    cur = conn.cursor()
    cur.execute("SET statement_timeout = '600s'")  # 10 minutes for large queries
    cur.close()
    print(f"  DEBUG: Connected successfully", flush=True)
    return conn


def _bytes_to_vector(data: bytes, seq_len: int, hidden_dim: int, strategy: str, prompt_len: int) -> List[float]:
    """Convert bytes to activation vector using extraction strategy."""
    num_floats = seq_len * hidden_dim
    values = struct.unpack(f'{num_floats}f', data)

    if strategy == "last_token":
        start = (seq_len - 1) * hidden_dim
        return list(values[start:start + hidden_dim])
    elif strategy == "first_token":
        idx = min(prompt_len, seq_len - 1)
        start = idx * hidden_dim
        return list(values[start:start + hidden_dim])
    else:
        start = (seq_len - 1) * hidden_dim
        return list(values[start:start + hidden_dim])


def load_activations_from_database(
    model_name: str,
    task_name: str,
    layer: int,
    prompt_format: str = "chat",
    extraction_strategy: str = "last_token",
    limit: Optional[int] = None,
    database_url: Optional[str] = None,
    pair_ids: Optional[set] = None,
    use_cache: bool = True,
    force_refresh: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Load activations for a single layer from local cache or database.

    Args:
        model_name: HuggingFace model ID (e.g., "meta-llama/Llama-3.2-1B-Instruct")
        task_name: Name of the contrastive pair set (e.g., "truthfulqa_custom")
        layer: Layer number to load
        prompt_format: Prompt format used ("chat" or "completion")
        extraction_strategy: Token extraction strategy ("last_token" or "first_token")
        limit: Maximum number of pairs to load
        database_url: Optional database URL (defaults to DATABASE_URL env var)
        pair_ids: Optional set of pair IDs to filter (for train/test split)
        use_cache: Whether to use local cache (default: True)
        force_refresh: Force reload from database even if cache exists

    Returns:
        Tuple of (pos_activations, neg_activations) tensors with pair_id mapping
    """
    cache_path = get_cache_path(task_name, "activations", model_name=model_name, layer=layer)

    # Try loading from cache first (only if no pair_ids filter)
    if use_cache and not force_refresh and pair_ids is None and cache_path.exists():
        print(f"  Loading activations from cache: {cache_path}")
        cached = torch.load(cache_path, weights_only=True)
        pos_tensor = cached["pos"]
        neg_tensor = cached["neg"]
        cached_pair_ids = cached.get("pair_ids", list(range(len(pos_tensor))))

        # Apply limit if specified
        if limit and len(pos_tensor) > limit:
            pos_tensor = pos_tensor[:limit]
            neg_tensor = neg_tensor[:limit]

        return pos_tensor, neg_tensor

    # If we have pair_ids filter, try to load from cache and filter
    if use_cache and not force_refresh and pair_ids is not None and cache_path.exists():
        print(f"  Loading activations from cache and filtering: {cache_path}")
        cached = torch.load(cache_path, weights_only=True)
        cached_pair_ids = cached.get("pair_ids", [])
        pos_all = cached["pos"]
        neg_all = cached["neg"]

        # Filter by pair_ids
        pos_list, neg_list = [], []
        for i, pid in enumerate(cached_pair_ids):
            if pid in pair_ids:
                pos_list.append(pos_all[i])
                neg_list.append(neg_all[i])

        if pos_list:
            return torch.stack(pos_list), torch.stack(neg_list)
        # If no matches, fall through to database
    from collections import defaultdict

    conn = _get_db_connection(database_url)
    cur = conn.cursor()

    try:
        # Get model ID
        cur.execute('SELECT id FROM "Model" WHERE "huggingFaceId" = %s', (model_name,))
        result = cur.fetchone()
        if not result:
            raise ValueError(f"Model {model_name} not found in database")
        model_id = result[0]

        # Get set ID
        cur.execute('SELECT id FROM "ContrastivePairSet" WHERE name = %s', (task_name,))
        result = cur.fetchone()
        if not result:
            raise ValueError(f"Benchmark {task_name} not found in database")
        set_id = result[0]

        # Load all activations for this layer
        limit_clause = f"LIMIT {limit * 2}" if limit else ""

        cur.execute(f'''
            SELECT "contrastivePairId", "seqLen", "hiddenDim", "promptLen", "hiddenStates", "isPositive"
            FROM "RawActivation"
            WHERE "modelId" = %s
              AND "contrastivePairSetId" = %s
              AND "promptFormat" = %s
              AND layer = %s
            ORDER BY "contrastivePairId", "isPositive"
            {limit_clause}
        ''', (model_id, set_id, prompt_format, layer))

        rows = cur.fetchall()

        # Group by pair_id
        pair_activations = defaultdict(dict)
        for row in rows:
            pair_id, seq_len, h_dim, prompt_len, hidden_bytes, is_positive = row
            # Filter by pair_ids if provided
            if pair_ids is not None and pair_id not in pair_ids:
                continue
            activation = _bytes_to_vector(hidden_bytes, seq_len, h_dim, extraction_strategy, prompt_len or 0)
            key = "pos" if is_positive else "neg"
            pair_activations[pair_id][key] = activation

        # Convert to tensors - only pairs with both pos and neg
        pos_list, neg_list, loaded_pair_ids = [], [], []
        for pair_id, acts in sorted(pair_activations.items()):
            if "pos" in acts and "neg" in acts:
                pos_list.append(acts["pos"])
                neg_list.append(acts["neg"])
                loaded_pair_ids.append(pair_id)

        pos_tensor = torch.tensor(pos_list, dtype=torch.float32)
        neg_tensor = torch.tensor(neg_list, dtype=torch.float32)

        # Save to cache if enabled and no pair_ids filter (cache full data)
        if use_cache and pair_ids is None and len(pos_list) > 0:
            print(f"  Caching {len(pos_list)} activations to: {cache_path}")
            torch.save({
                "pos": pos_tensor,
                "neg": neg_tensor,
                "pair_ids": loaded_pair_ids,
            }, cache_path)

        return pos_tensor, neg_tensor

    finally:
        cur.close()
        conn.close()


def load_pair_texts_from_database(
    task_name: str,
    limit: int = 200,
    database_url: Optional[str] = None,
    use_cache: bool = True,
    force_refresh: bool = False,
) -> Dict[int, Dict[str, str]]:
    """
    Load contrastive pair texts from local cache or Supabase database.

    Args:
        task_name: Name of the contrastive pair set (e.g., "truthfulqa_custom")
        limit: Maximum number of pairs to load
        database_url: Optional database URL (defaults to DATABASE_URL env var)
        use_cache: Whether to use local cache (default: True)
        force_refresh: Force reload from database even if cache exists

    Returns:
        Dict mapping pair_id -> {prompt, positive, negative}
    """
    print(f"  DEBUG: load_pair_texts_from_database called", flush=True)
    cache_path = get_cache_path(task_name, "pair_texts")
    print(f"  DEBUG: cache_path = {cache_path}", flush=True)
    print(f"  DEBUG: cache_path.exists() = {cache_path.exists()}", flush=True)
    print(f"  DEBUG: use_cache={use_cache}, force_refresh={force_refresh}", flush=True)

    # Try loading from cache first
    if use_cache and not force_refresh and cache_path.exists():
        print(f"  Loading pair texts from cache: {cache_path}", flush=True)
        with open(cache_path, 'r') as f:
            cached_data = json.load(f)
        print(f"  DEBUG: loaded {len(cached_data)} pairs from cache", flush=True)
        # Convert string keys back to int
        pairs = {int(k): v for k, v in cached_data.items()}
        # Apply limit
        if limit and len(pairs) > limit:
            pair_ids = sorted(pairs.keys())[:limit]
            pairs = {pid: pairs[pid] for pid in pair_ids}
        print(f"  DEBUG: returning {len(pairs)} pairs", flush=True)
        return pairs

    # Load from database
    conn = _get_db_connection(database_url)
    cur = conn.cursor()

    try:
        # Get set ID for task
        cur.execute('SELECT id FROM "ContrastivePairSet" WHERE name = %s', (task_name,))
        result = cur.fetchone()
        if not result:
            raise ValueError(f"No ContrastivePairSet found with name '{task_name}'")
        set_id = result[0]

        # Load pairs (load all for caching, apply limit after)
        cur.execute('''
            SELECT id, "positiveExample", "negativeExample"
            FROM "ContrastivePair"
            WHERE "setId" = %s
            ORDER BY id
        ''', (set_id,))

        pairs = {}
        for row in cur.fetchall():
            pair_id, pos_example, neg_example = row

            # Parse prompt and responses
            if "\n\n" in pos_example:
                prompt = pos_example.rsplit("\n\n", 1)[0]
                pos_response = pos_example.rsplit("\n\n", 1)[1]
            else:
                prompt = pos_example
                pos_response = ""

            if "\n\n" in neg_example:
                neg_response = neg_example.rsplit("\n\n", 1)[1]
            else:
                neg_response = neg_example

            pairs[pair_id] = {
                "prompt": prompt,
                "positive": pos_response,
                "negative": neg_response,
            }

        # Save to cache
        if use_cache:
            print(f"  Caching {len(pairs)} pair texts to: {cache_path}")
            with open(cache_path, 'w') as f:
                json.dump(pairs, f)

        # Apply limit
        if limit and len(pairs) > limit:
            pair_ids = sorted(pairs.keys())[:limit]
            pairs = {pid: pairs[pid] for pid in pair_ids}

        return pairs

    finally:
        cur.close()
        conn.close()
