"""Handles interaction with the Goodfire API for getting activations."""

import asyncio
import logging
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import cast

import goodfire
import httpx
import torch
from dotenv import load_dotenv
from goodfire.api.chat.interfaces import ChatMessage
from jaxtyping import Float
from torch import Tensor
from tqdm import tqdm

from deception_detection.log import logger
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Dialogue

# Thread-local storage for Goodfire clients
thread_local = threading.local()


def get_thread_local_client() -> goodfire.Client:
    if not hasattr(thread_local, "client"):
        load_dotenv()
        api_key = os.getenv("GOODFIRE_API_KEY")
        assert api_key is not None
        thread_local.client = goodfire.Client(api_key)
    return thread_local.client


def get_goodfire_activations(
    tokenized_dataset: TokenizedDataset, model_size: int = 70
) -> Float[Tensor, "batch seqpos layer emb"]:
    """Get activations from Goodfire API for a tokenized dataset.

    Returns:
        Sparse tensor of shape (batch, seqpos, 1, emb)
    """
    match model_size:
        case 70:
            variant = goodfire.Variant("meta-llama/Llama-3.3-70B-Instruct")
        case 8:
            variant = goodfire.Variant("meta-llama/Meta-Llama-3.1-8B-Instruct")
        case _:
            raise ValueError(f"Unsupported model size: {model_size}")

    # Calculate max sequence length from tokenized dataset
    max_seq_len = tokenized_dataset.tokens.shape[1]

    # Get all contexts in parallel while preserving order
    all_sparse_mats_placeholders: list[Float[Tensor, "seqpos layer emb"] | None] = [None] * len(
        tokenized_dataset.dialogues
    )
    with ThreadPoolExecutor(max_workers=40) as executor:
        future_to_idx = {
            executor.submit(
                get_one_matrix, dialogue, variant, tokenized_dataset.str_tokens[idx], max_seq_len
            ): idx
            for idx, dialogue in enumerate(tokenized_dataset.dialogues)
        }

        with tqdm(
            total=len(tokenized_dataset.dialogues),
            desc=f"Goodfire - {tokenized_dataset.dataset_partial_id}",
        ) as pbar:
            for future in as_completed(future_to_idx):
                idx = future_to_idx[future]
                try:
                    sparse_mat = future.result()
                    all_sparse_mats_placeholders[idx] = sparse_mat
                    pbar.update(1)
                except Exception as e:
                    logger.error(f"Error getting matrix for dialogue {idx}: {e}")
                    raise e

    # Verify we got all matrices
    assert all([m is not None for m in all_sparse_mats_placeholders])
    all_sparse_mats: list[Float[Tensor, "seqpos layer emb"]] = cast(
        list[Float[Tensor, "seqpos layer emb"]], all_sparse_mats_placeholders
    )

    return torch.stack(all_sparse_mats, dim=0)


def get_one_matrix(
    dialogue: Dialogue, variant: goodfire.Variant, expected_tokens: list[str], max_seq_len: int
) -> Float[Tensor, "seqpos layer emb"]:
    """Get context from Goodfire API for a single dialogue and convert to padded sparse matrix."""
    client = get_thread_local_client()

    # Disable noisy HTTP info logs
    logging.getLogger("httpx").setLevel(logging.WARNING)

    messages: list[ChatMessage] = [
        {"role": message.role, "content": message.content} for message in dialogue
    ]

    # merge adjacent messages with the same role
    merged_messages: list[ChatMessage] = []
    for message in messages:
        if merged_messages and merged_messages[-1]["role"] == message["role"]:
            merged_messages[-1]["content"] += message["content"]
        else:
            merged_messages.append(message)

    context = None
    for attempt in range(3):
        try:
            # Create new event loop for this thread if it doesn't exist
            try:
                loop = asyncio.get_event_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)

            context = client.features.inspect(
                merged_messages, model=variant, _fetch_feature_data=False
            )
            break
        except httpx.ReadTimeout:
            if attempt == 2:  # Last attempt
                raise
            logger.warning(f"Timeout on attempt {attempt+1}/3, retrying...")
            time.sleep(1 * (attempt + 1))  # Exponential backoff

    assert context is not None  # For type checker

    # Verify tokens match between dataset and Goodfire context
    goodfire_tokens = [t._token for t in context.tokens]  # type: ignore
    if goodfire_tokens != expected_tokens:
        logger.error("Token mismatch:")
        logger.error(f"Goodfire: {goodfire_tokens}")
        logger.error(f"Dataset: {expected_tokens}")
        raise ValueError("Token mismatch between Goodfire and dataset")

    # Convert to padded sparse matrix
    mat = torch.tensor(context.matrix())
    mat = torch.cat([torch.zeros(max_seq_len - mat.shape[0], mat.shape[1]), mat], dim=0)
    mat = mat.unsqueeze(1)  # add layer dim (always index 0)
    mat = mat.to_sparse_coo()  # make sparse
    return mat
