from typing import List, Dict, Any
from pydantic import BaseModel
import pandas as pd


class PromptLog(BaseModel):
    prompt_id: str
    data: Dict[str, Any]


class PromptLogs(BaseModel):
    prompt_logs: Dict[str, List[PromptLog]]


def build_df_from_prompt_logs(prompt_logs: PromptLogs) -> Dict[str, pd.DataFrame]:
    dfs = {}
    for df_name, logs in prompt_logs.prompt_logs.items():
        data = []
        for log in logs:
            data.append({"prompt_id": log.prompt_id, **log.data})
        if not data:
            print(f"No data for {df_name}, skipping")
            continue
        df = pd.DataFrame(data).set_index("prompt_id")
        df.name = df_name
        dfs[df_name] = df
    if not dfs:
        return
    return pd.concat(dfs.values(), axis=1).reset_index(drop=True)
