import json
import os

from minio import Minio

mc = Minio(os.environ["MINIO_ENDPOINT"], access_key=os.environ["MINIO_ACCESS_KEY"], secret_key=os.environ["MINIO_SECRET_KEY"], secure=False)
raw_objs = list(mc.list_objects("stonks-llm-results", recursive=True))

# Count companies found vs not found
total = 0
with_companies = 0
empty_companies = 0
wrong_tickers = 0
our_tickers = {"AAPL","MSFT","NVDA","AMZN","GOOGL","JPM","JNJ","XOM","TSLA","META"}
ticker_hits = {}

for o in raw_objs:
    data = json.loads(mc.get_object("stonks-llm-results", o.object_name).read())
    if not data.get("success"):
        continue
    attempts = data.get("attempts", [])
    if not attempts:
        continue
    raw_out = attempts[-1].get("raw_output", "")
    try:
        parsed = json.loads(raw_out)
        companies = parsed.get("companies", [])
        total += 1
        if companies:
            with_companies += 1
            for c in companies:
                t = c.get("ticker", "")
                if t in our_tickers:
                    ticker_hits[t] = ticker_hits.get(t, 0) + 1
                else:
                    wrong_tickers += 1
        else:
            empty_companies += 1
    except:
        pass

print(f"Total successful extractions: {total}")
print(f"  With companies: {with_companies}")
print(f"  Empty companies: {empty_companies}")
print(f"  Wrong/unknown tickers: {wrong_tickers}")
print(f"  Our ticker hits: {ticker_hits}")
