import json
import os
import asyncio
import time
from typing import Any
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from openai import AsyncOpenAI
import anthropic
from pydantic import BaseModel
import tqdm
from tqdm.asyncio import tqdm as async_tqdm

def default_chat_template(messages: list[dict[str, str]], tokenize: bool = False, add_generation_prompt: bool = True) -> str:
    prompt = ""
    for message in messages:
        prompt += f"{message['role']}: {message['content']}\n"
    if add_generation_prompt:
        prompt += "assistant: "
    return prompt


def generate_responses(
    model_name: str,
    prompts: list[str],
    system_prompt: str = "",
    max_new_tokens: int = 256,
    batch_size: int = 4,
    provider: str = "hf",  # "hf" | "openai" | "anthropic"
    temperature: float = 1.0
):
    if provider == "hf":
        return generate_hf_responses(
            model_name, prompts, system_prompt, max_new_tokens, batch_size
        )
    elif provider == "openai":
        try:
            # Check if there's already an event loop running
            try:
                loop = asyncio.get_running_loop()
                # If we're already in an event loop, we can't use asyncio.run()
                # This means we're already in an async context
                raise RuntimeError("Cannot use asyncio.run() in an existing event loop")
            except RuntimeError:
                # No event loop running, safe to use asyncio.run()
                return asyncio.run(generate_openai_responses(
                    model_name, prompts, system_prompt, max_new_tokens, temperature
                ))
        except Exception as e:
            print(f"Error in OpenAI API call: {e}")
            # Return empty responses on error
            return [""] * len(prompts)
    elif provider == "anthropic":
        return generate_anthropic_responses(
            model_name, prompts, system_prompt, max_new_tokens
        )
    else:
        raise ValueError(
            "Unsupported provider: choose 'hf', 'openai', or 'anthropic'")


def generate_hf_responses(
    model_name: str,
    prompts: list[str],
    system_prompt: str,
    max_new_tokens: int,
    batch_size: int,
) -> list[str]:
    tokenizer = AutoTokenizer.from_pretrained(
        model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        low_cpu_mem_usage=True,
    )
    model.eval()

    results = []
    for i in range(0, len(prompts), batch_size):
        batch = prompts[i: i + batch_size]
        messages_batch = []
        for user_prompt in batch:
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            messages.append({"role": "user", "content": user_prompt})
            if (
                hasattr(tokenizer, "apply_chat_template")
                and callable(getattr(tokenizer, "apply_chat_template", None))
                and getattr(tokenizer, "chat_template", None) is not None
            ):
                chat_input = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
            else:
                chat_input = default_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
            messages_batch.append(chat_input)

        tokenized = tokenizer(
            messages_batch,
            return_tensors="pt",
            padding=True,
            truncation=True,
        ).to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **tokenized,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )

        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        for inp, out in zip(messages_batch, decoded, strict=False):
            response = out[len(inp):].strip()
            results.append(response)

    return results


async def generate_openai_responses(model_name, prompts, system_prompt, max_tokens, temperature):
    client = AsyncOpenAI()  # uses OPENAI_API_KEY from environment

    async def fetch_response(prompt):
        try:
            response = await client.chat.completions.create(
                model=model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=max_tokens,
                temperature=temperature,
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"Error fetching response: {e}")
            return ""

    try:
        # Launch all requests concurrently with progress bar
        tasks = [fetch_response(prompt) for prompt in prompts]
        responses = await async_tqdm.gather(*tasks, desc="Generating responses")
        return responses
    finally:
        # Ensure client is properly closed
        try:
            await client.aclose()
        except Exception as e:
            # Ignore cleanup errors
            pass


def generate_anthropic_responses(model_name: str, prompts: list[str], system_prompt: str, max_tokens: int) -> list[str]:
    client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
    responses = []

    for prompt in prompts:
        message = client.messages.create(
            model=model_name,
            max_tokens=max_tokens,
            # type: ignore[arg-type]
            system=system_prompt if system_prompt else None,
            messages=[{"role": "user", "content": prompt}],
        )
        # type: ignore[union-attr]
        responses.append(message.content[0].text.strip())

    return responses


async def generate_structured_response(format: str, model_name: str, system_prompt: str) -> list[dict[str, str]]:
    client = AsyncOpenAI()  # uses OPENAI_API_KEY from environment
    t0 = time.time()
    response = await client.responses.parse(
        model=model_name,
        input=[
            {"role": "system", "content": system_prompt},
        ],
        text_format=format,  # type: ignore[arg-type]
        max_output_tokens=10000,
    )
    print(f"Response from Openai in {time.time() - t0:.2f}s")
    return response.output_parsed  # type: ignore[return-value]


async def generate_structured_responses(
    format: str,
    model_name: str,
    system_prompt: str,
    num_responses: int = 3,
) -> list[list[dict[str, str]]]:
    api_calls = [
        generate_structured_response(format, model_name, system_prompt)
        for _ in range(num_responses)
    ]
    responses = await asyncio.gather(*api_calls, return_exceptions=True)
    return [r for r in responses if r is not None]  # type: ignore[misc]


if __name__ == "__main__":
    from pydantic import BaseModel

    class TestFormat(BaseModel):
        name: str
        path: str

    async def wrapper() -> None:
        res = await generate_structured_responses(
            TestFormat,  # type: ignore[arg-type]
            "gpt-4o-mini",
            "Generate a random filepath and filename",
            num_responses=3,
        )
        print(res)

    asyncio.run(wrapper())
