#!/usr/bin/env python3
"""
Model tracker for agata chat command.

Tracks successful and unsuccessful model usage with timestamps.
Stores configuration in ~/.config/agata/
"""

import os
import yaml
from datetime import datetime
from typing import Dict, Optional, List
from pathlib import Path

CONFIG_DIR = Path.home() / ".config" / "agata"
FAILED_MODELS_FILE = CONFIG_DIR / "failed_models.yaml"
SUCCESSFUL_MODELS_FILE = CONFIG_DIR / "successful_models.yaml"


def _ensure_config_dir():
    """Ensure the config directory exists."""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)


def _load_yaml(filepath: Path) -> Dict:
    """Load YAML file or return empty dict if not exists."""
    if not filepath.exists():
        return {}
    try:
        with open(filepath, "r") as f:
            return yaml.safe_load(f) or {}
    except Exception:
        return {}


def _save_yaml(filepath: Path, data: Dict):
    """Save data to YAML file."""
    _ensure_config_dir()
    with open(filepath, "w") as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=True)


def get_failed_models() -> Dict[str, str]:
    """
    Get dictionary of failed models with their failure timestamps.

    Returns:
        Dict mapping model names to ISO format timestamps
    """
    return _load_yaml(FAILED_MODELS_FILE)


def get_successful_models() -> Dict[str, str]:
    """
    Get dictionary of successful models with their last use timestamps.

    Returns:
        Dict mapping model names to ISO format timestamps
    """
    return _load_yaml(SUCCESSFUL_MODELS_FILE)


def add_failed_model(model: str, timestamp: Optional[str] = None):
    """
    Add a model to the failed models list.
    Removes it from successful models if present.

    Args:
        model: The model identifier (e.g., 'openai/gpt-4')
        timestamp: ISO format timestamp (defaults to current time)
    """
    if not model:
        return

    if timestamp is None:
        timestamp = datetime.now().isoformat()

    # Remove from successful models if present
    successful = get_successful_models()
    if model in successful:
        del successful[model]
        _save_yaml(SUCCESSFUL_MODELS_FILE, successful)

    # Add to failed models
    failed = get_failed_models()
    failed[model] = timestamp
    _save_yaml(FAILED_MODELS_FILE, failed)


def add_successful_model(model: str, timestamp: Optional[str] = None):
    """
    Add a model to the successful models list.
    Removes it from failed models if present.

    Args:
        model: The model identifier (e.g., 'openai/gpt-4')
        timestamp: ISO format timestamp (defaults to current time)
    """
    if not model:
        return

    if timestamp is None:
        timestamp = datetime.now().isoformat()

    # Remove from failed models if present
    failed = get_failed_models()
    if model in failed:
        del failed[model]
        _save_yaml(FAILED_MODELS_FILE, failed)

    # Add to successful models
    successful = get_successful_models()
    successful[model] = timestamp
    _save_yaml(SUCCESSFUL_MODELS_FILE, successful)


def is_rate_limit_error(error_message: str) -> bool:
    """
    Check if an error message indicates a rate limit (429) error.

    Args:
        error_message: The error message string

    Returns:
        True if it's a 429/Too Many Requests error
    """
    if not error_message:
        return False

    error_lower = error_message.lower()
    rate_limit_indicators = ["429", "too many requests", "rate limit", "ratelimit"]

    return any(indicator in error_lower for indicator in rate_limit_indicators)


def track_model_usage(model: str, success: bool, error_message: Optional[str] = None):
    """
    Track model usage based on success or failure.

    Tracks all failed models regardless of error type.

    Args:
        model: The model identifier
        success: Whether the request was successful
        error_message: Error message if failed
    """
    if not model:
        return

    if success:
        add_successful_model(model)
    else:
        add_failed_model(model)


def list_models() -> Dict[str, Dict[str, str]]:
    """
    Get both failed and successful model lists.

    Returns:
        Dict with 'failed' and 'successful' keys containing model dicts
    """
    return {"failed": get_failed_models(), "successful": get_successful_models()}


def clear_tracking():
    """Clear all model tracking data."""
    _ensure_config_dir()
    if FAILED_MODELS_FILE.exists():
        FAILED_MODELS_FILE.unlink()
    if SUCCESSFUL_MODELS_FILE.exists():
        SUCCESSFUL_MODELS_FILE.unlink()
