import json
import os
from datetime import datetime
from typing import Dict, Optional
from pr.core.logging import get_logger

logger = get_logger('usage')

USAGE_DB_FILE = os.path.expanduser("~/.assistant_usage.json")

MODEL_COSTS = {
    'x-ai/grok-code-fast-1': {'input': 0.0, 'output': 0.0},
    'gpt-4': {'input': 0.03, 'output': 0.06},
    'gpt-4-turbo': {'input': 0.01, 'output': 0.03},
    'gpt-3.5-turbo': {'input': 0.0005, 'output': 0.0015},
    'claude-3-opus': {'input': 0.015, 'output': 0.075},
    'claude-3-sonnet': {'input': 0.003, 'output': 0.015},
    'claude-3-haiku': {'input': 0.00025, 'output': 0.00125},
}


class UsageTracker:

    def __init__(self):
        self.session_usage = {
            'requests': 0,
            'total_tokens': 0,
            'input_tokens': 0,
            'output_tokens': 0,
            'estimated_cost': 0.0,
            'models_used': {}
        }

    def track_request(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        total_tokens: Optional[int] = None
    ):
        if total_tokens is None:
            total_tokens = input_tokens + output_tokens

        self.session_usage['requests'] += 1
        self.session_usage['total_tokens'] += total_tokens
        self.session_usage['input_tokens'] += input_tokens
        self.session_usage['output_tokens'] += output_tokens

        if model not in self.session_usage['models_used']:
            self.session_usage['models_used'][model] = {
                'requests': 0,
                'tokens': 0,
                'cost': 0.0
            }

        model_usage = self.session_usage['models_used'][model]
        model_usage['requests'] += 1
        model_usage['tokens'] += total_tokens

        cost = self._calculate_cost(model, input_tokens, output_tokens)
        model_usage['cost'] += cost
        self.session_usage['estimated_cost'] += cost

        self._save_to_history(model, input_tokens, output_tokens, cost)

        logger.debug(
            f"Tracked request: {model}, tokens: {total_tokens}, cost: ${cost:.4f}"
        )

    def _calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        if model not in MODEL_COSTS:
            base_model = model.split('/')[0] if '/' in model else model
            if base_model not in MODEL_COSTS:
                logger.warning(f"Unknown model for cost calculation: {model}")
                return 0.0
            costs = MODEL_COSTS[base_model]
        else:
            costs = MODEL_COSTS[model]

        input_cost = (input_tokens / 1000) * costs['input']
        output_cost = (output_tokens / 1000) * costs['output']

        return input_cost + output_cost

    def _save_to_history(self, model: str, input_tokens: int, output_tokens: int, cost: float):
        try:
            history = []
            if os.path.exists(USAGE_DB_FILE):
                with open(USAGE_DB_FILE, 'r') as f:
                    history = json.load(f)

            history.append({
                'timestamp': datetime.now().isoformat(),
                'model': model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'total_tokens': input_tokens + output_tokens,
                'cost': cost
            })

            if len(history) > 10000:
                history = history[-10000:]

            with open(USAGE_DB_FILE, 'w') as f:
                json.dump(history, f, indent=2)

        except Exception as e:
            logger.error(f"Error saving usage history: {e}")

    def get_session_summary(self) -> Dict:
        return self.session_usage.copy()

    def get_formatted_summary(self) -> str:
        usage = self.session_usage
        lines = [
            "\n=== Session Usage Summary ===",
            f"Total Requests: {usage['requests']}",
            f"Total Tokens: {usage['total_tokens']:,}",
            f"  Input: {usage['input_tokens']:,}",
            f"  Output: {usage['output_tokens']:,}",
            f"Estimated Cost: ${usage['estimated_cost']:.4f}",
        ]

        if usage['models_used']:
            lines.append("\nModels Used:")
            for model, stats in usage['models_used'].items():
                lines.append(
                    f"  {model}: {stats['requests']} requests, "
                    f"{stats['tokens']:,} tokens, ${stats['cost']:.4f}"
                )

        return '\n'.join(lines)

    @staticmethod
    def get_total_usage() -> Dict:
        if not os.path.exists(USAGE_DB_FILE):
            return {
                'total_requests': 0,
                'total_tokens': 0,
                'total_cost': 0.0
            }

        try:
            with open(USAGE_DB_FILE, 'r') as f:
                history = json.load(f)

            total_tokens = sum(entry['total_tokens'] for entry in history)
            total_cost = sum(entry['cost'] for entry in history)

            return {
                'total_requests': len(history),
                'total_tokens': total_tokens,
                'total_cost': total_cost
            }

        except Exception as e:
            logger.error(f"Error loading usage history: {e}")
            return {
                'total_requests': 0,
                'total_tokens': 0,
                'total_cost': 0.0
            }