import threading

_lock = threading.Lock()
_session_tokens = {}


def init_session(chat_id: str):
    with _lock:
        _session_tokens[chat_id] = {
            "classify":   {"input": 0, "output": 0, "calls": 0},
            "extract":    {"input": 0, "output": 0, "calls": 0},
            "validate":   {"input": 0, "output": 0, "calls": 0},
            "exceptions": {"input": 0, "output": 0, "calls": 0},
            "rca":        {"input": 0, "output": 0, "calls": 0},
            "capa":       {"input": 0, "output": 0, "calls": 0},
            "report":     {"input": 0, "output": 0, "calls": 0},
        }


def add_tokens(chat_id: str, step: str, input_tokens: int, output_tokens: int):
    with _lock:
        if chat_id not in _session_tokens:
            init_session(chat_id)
        if step in _session_tokens[chat_id]:
            _session_tokens[chat_id][step]["input"]  += input_tokens
            _session_tokens[chat_id][step]["output"] += output_tokens
            _session_tokens[chat_id][step]["calls"]  += 1
            print(f"[TOKENS] {step} — input:{input_tokens} output:{output_tokens} total:{input_tokens + output_tokens}")


def get_summary(chat_id: str) -> dict:
    with _lock:
        if chat_id not in _session_tokens:
            return {}
        steps = _session_tokens[chat_id]
        summary = {}
        grand_input = grand_output = grand_calls = 0

        for step, data in steps.items():
            total = data["input"] + data["output"]
            summary[step] = {
                "input_tokens":  data["input"],
                "output_tokens": data["output"],
                "total_tokens":  total,
                "llm_calls":     data["calls"]
            }
            grand_input  += data["input"]
            grand_output += data["output"]
            grand_calls  += data["calls"]

        summary["TOTAL"] = {
            "input_tokens":  grand_input,
            "output_tokens": grand_output,
            "total_tokens":  grand_input + grand_output,
            "llm_calls":     grand_calls,
            "estimated_cost_usd": round((grand_input * 0.000005) + (grand_output * 0.000015), 4)
        }
        return summary


def clear_session(chat_id: str):
    with _lock:
        _session_tokens.pop(chat_id, None)