#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 15 14:12:21 2025

@author: samd
"""
import requests
import json


def or_costs(model_name, ninput_tokens, noutput_tokens):
    costs = or_token_costs(model_name)
    input_cost = float(costs['prompt_cost'])
    output_cost = float(costs['completion_cost'])
    total_cost = ninput_tokens*input_cost + noutput_tokens*output_cost
    
    return total_cost

def or_token_costs(model_name=None):
    """
    Extract costs for OpenRouter models
    
    Args:
        model_name (str, optional): Specific model to get costs for. 
                                   If None, returns all models.
    
    Returns:
        dict: Model costs information
    """
    
    url = "https://openrouter.ai/api/v1/models"
    
    try:
        response = requests.get(url)
        response.raise_for_status()
        
        models_data = response.json()
        
        if model_name:
            # Find specific model
            for model in models_data.get('data', []):
                if model['id'] == model_name:
                    return extract_model_info(model)
            return f"Model '{model_name}' not found"
        else:
            # Return all models with costs
            all_models = {}
            for model in models_data.get('data', []):
                all_models[model['id']] = extract_model_info(model)
            return all_models
            
    except requests.RequestException as e:
        return f"Error fetching data: {e}"

def extract_model_info(model):
    """Extract relevant cost and model information"""
    return {
        'name': model.get('id', 'N/A'),
        'prompt_cost': model.get('pricing', {}).get('prompt', 'N/A'),
        'completion_cost': model.get('pricing', {}).get('completion', 'N/A'),
        'context_length': model.get('context_length', 'N/A'),
        'description': model.get('description', 'N/A')
    }