from langchain_openai import ChatOpenAI, AzureChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
import json, os

GLOBAL_LLM = None

def get_llm(log, llm_configuration=None, verbose=True):
    log.write("Initiating connection to Large Language Model API...", verbose=verbose)
    if llm_configuration is None:
        key_path = os.path.expanduser("~/.gwaslab/LLM_KEY")
        with open(key_path, "r", encoding="utf-8") as f:
            llm_configuration = json.load(f)
    
    global GLOBAL_LLM
    
    if GLOBAL_LLM is None:
        provider = llm_configuration.get("provider", "openai")
        log.write(" -Connecting to Large Language Model API...", verbose=verbose)
        if provider == "openai":
            config = {k: v for k, v in llm_configuration.items() if k != "provider"}
            GLOBAL_LLM = ChatOpenAI(**config)
            log.write(" -Detected OpenAI-like API...", verbose=verbose)
        elif provider == "google":
            config = {k: v for k, v in llm_configuration.items() if k != "provider"}
            GLOBAL_LLM = ChatGoogleGenerativeAI(**config)
            log.write(" -Detected Google API...", verbose=verbose)
        elif provider == "azure":
            config = {k: v for k, v in llm_configuration.items() if k != "provider"}
            GLOBAL_LLM = AzureChatOpenAI(**config)
            log.write(" -Detected Azure OpenAI API...", verbose=verbose)
        else:
            raise ValueError(f"Unsupported LLM provider: {provider}")
    model_name = (
        getattr(GLOBAL_LLM, "model", None)
        or getattr(GLOBAL_LLM, "model_name", None)
    )
    if model_name is not None:
        log.write(f" -Model used: {model_name}...", verbose=verbose)
    return GLOBAL_LLM

# ---------------------------------------------------------------------------------
# Token usage helpers
# ---------------------------------------------------------------------------------

def extract_token_usage(msg):
    usage = getattr(msg, "response_metadata", {}).get("token_usage") \
            or getattr(msg, "usage_metadata", None)
    if not usage:
        return None
    prompt = usage.get("prompt_tokens", usage.get("input_tokens"))
    completion = usage.get("completion_tokens", usage.get("output_tokens"))
    total = usage.get("total_tokens")
    return {
        "input": prompt if isinstance(prompt, int) else 0,
        "output": completion if isinstance(completion, int) else 0,
        "total": total if isinstance(total, int) else None,
    }

def accumulate_token_usage(counter: dict, usage: dict):
    if not isinstance(counter, dict) or not isinstance(usage, dict):
        return
    p = usage.get("input", 0) or 0
    c = usage.get("output", 0) or 0
    t = usage.get("total")
    if not isinstance(t, int):
        t = p + c
    counter["input"] = counter.get("input", 0) + p
    counter["output"] = counter.get("output", 0) + c
    counter["total"] = counter.get("total", 0) + t

def snapshot_counters(counter: dict):
    return (
        counter.get("input", 0),
        counter.get("output", 0),
        counter.get("total", 0),
    )

def log_run_totals(log, label: str, start_tuple, end_tuple, verbose=True):
    si, so, st = start_tuple
    ei, eo, et = end_tuple
    di = max(0, ei - si)
    do = max(0, eo - so)
    dt = max(0, et - st)
    log.write(f"[USAGE] This {label}: prompt={di}, completion={do}, total={dt}", verbose=verbose)
    log.write(f"[USAGE] Accumulative: prompt={ei}, completion={eo}, total={et}", verbose=verbose)
