"""Estimate GraphRouter dataset generation cost across the LLM pool.

The estimator assumes that every query in the unified QA dataset will be sent
once to each configured LLM together with a maximum output length equal to the
`query_response_length` field in `configs/config.yaml`. It relies on the
`tiktoken` `cl100k_base` tokenizer to approximate prompt token counts for all
models, so treat the results as an upper bound rather than an exact bill.
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable

import pandas as pd
import tiktoken
import yaml

ENCODING = tiktoken.get_encoding("cl100k_base")


def load_config(path: Path) -> Dict[str, object]:
    with path.open("r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def load_llm_descriptions(path: Path) -> Dict[str, Dict[str, object]]:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def load_dataset(path: Path) -> pd.DataFrame:
    return pd.read_csv(path)


def estimate_cost(
    dataset: pd.DataFrame,
    llm_description: Dict[str, Dict[str, object]],
    output_tokens: int,
) -> Dict[str, Dict[str, float]]:
    per_llm_cost: Dict[str, float] = defaultdict(float)
    per_task_cost: Dict[str, float] = defaultdict(float)

    for _, row in dataset.iterrows():
        prompt = row["query"]
        task_id = row["task_id"]
        input_tokens = len(ENCODING.encode(prompt))

        for name, desc in llm_description.items():
            input_price = float(desc["input_price"])
            output_price = float(desc["output_price"])
            cost = (input_tokens / 1_000_000) * input_price
            cost += (output_tokens / 1_000_000) * output_price

            per_llm_cost[name] += cost
            per_task_cost[task_id] += cost

    total_cost = sum(per_llm_cost.values())

    return {
        "per_llm": dict(sorted(per_llm_cost.items(), key=lambda kv: kv[1], reverse=True)),
        "per_task": dict(sorted(per_task_cost.items(), key=lambda kv: kv[1], reverse=True)),
        "total": total_cost,
    }


def format_currency(value: float) -> str:
    return f"${value:,.2f}"


def print_report(
    report: Dict[str, Dict[str, float]], total_queries: int, output_tokens: int
) -> None:
    print("=== GraphRouter Dataset Generation Cost Estimate ===")
    print(f"Queries: {total_queries}")
    print(f"LLM calls: {total_queries * len(report['per_llm'])}")
    print(f"Output tokens per call (assumed): {output_tokens}")
    print()

    print("Per-LLM cost (descending):")
    for name, cost in report["per_llm"].items():
        print(f"  {name:25s} {format_currency(cost)}")
    print()

    print("Per-task cost (descending):")
    for task, cost in report["per_task"].items():
        print(f"  {task:25s} {format_currency(cost)}")
    print()

    print(f"Estimated total: {format_currency(report['total'])}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--config",
        type=Path,
        default=Path("configs/config.yaml"),
        help="Path to the GraphRouter config file",
    )
    parser.add_argument(
        "--dataset",
        type=Path,
        default=None,
        help="Optional path to a pre-built unified QA CSV (defaults to config entry)",
    )
    parser.add_argument(
        "--llm-descriptions",
        type=Path,
        default=None,
        help="Optional path to LLM_Descriptions.json (defaults to config entry)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = load_config(args.config)

    dataset_path = (
        args.dataset if args.dataset is not None else Path(config["unified_qa_data_path"])
    )
    llm_path = (
        args.llm_descriptions
        if args.llm_descriptions is not None
        else Path(config["llm_description_path"])
    )

    dataset = load_dataset(dataset_path)
    llm_description = load_llm_descriptions(llm_path)
    output_tokens = int(config["query_response_length"])

    report = estimate_cost(dataset, llm_description, output_tokens)
    print_report(report, len(dataset), output_tokens)


if __name__ == "__main__":
    main()
