"""
Shared utility functions used across the project.
"""

from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional


# ---------------------------------------------------------------------------
# Logging setup
# ---------------------------------------------------------------------------

def setup_logging(level: str = "INFO") -> None:
    """Configure root logger with a clean format."""
    logging.basicConfig(
        level=getattr(logging, level.upper(), logging.INFO),
        format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
        datefmt="%H:%M:%S",
    )


# ---------------------------------------------------------------------------
# Environment helpers
# ---------------------------------------------------------------------------

def get_env(key: str, default: Optional[str] = None, required: bool = False) -> Optional[str]:
    """
    Retrieve an environment variable with optional validation.

    Args:
        key: Environment variable name
        default: Default value if not set
        required: Raise ValueError if not set and no default
    """
    value = os.getenv(key, default)
    if required and not value:
        raise ValueError(
            f"Required environment variable '{key}' is not set. "
            f"Check your .env file."
        )
    return value


# ---------------------------------------------------------------------------
# JSON helpers
# ---------------------------------------------------------------------------

def load_json(path: Path | str) -> Any:
    """Load and return JSON from a file."""
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"JSON file not found: {path}")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json(data: Any, path: Path | str, indent: int = 2) -> None:
    """Save data as JSON to a file, creating parent dirs as needed."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=indent, ensure_ascii=False)


# ---------------------------------------------------------------------------
# Text helpers
# ---------------------------------------------------------------------------

def truncate(text: str, max_chars: int = 200, suffix: str = "...") -> str:
    """Truncate text to max_chars, appending suffix if truncated."""
    if len(text) <= max_chars:
        return text
    return text[:max_chars - len(suffix)] + suffix


def count_words(text: str) -> int:
    return len(text.split())


# ---------------------------------------------------------------------------
# Config loader
# ---------------------------------------------------------------------------

def load_config(config_path: Path | str = ".env") -> Dict[str, str]:
    """
    Load key=value pairs from a .env file into a dict.
    Does NOT set environment variables — use python-dotenv for that.
    """
    config: Dict[str, str] = {}
    path = Path(config_path)
    if not path.exists():
        return config
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, _, value = line.partition("=")
                config[key.strip()] = value.strip().strip('"').strip("'")
    return config


# ---------------------------------------------------------------------------
# Eval dataset helpers
# ---------------------------------------------------------------------------

def load_eval_dataset(paths: List[Path | str]) -> List[dict]:
    """Load and merge multiple eval JSON files into a single list."""
    combined: List[dict] = []
    for path in paths:
        try:
            data = load_json(path)
            if isinstance(data, list):
                combined.extend(data)
            elif isinstance(data, dict) and "prompts" in data:
                combined.extend(data["prompts"])
        except Exception as e:
            logging.getLogger(__name__).warning("Could not load eval file %s: %s", path, e)
    return combined


def format_cost_table(
    oss_tokens: int,
    frontier_tokens: int,
    oss_cost_per_1k: float = 0.0,
    frontier_cost_per_1k: float = 0.0006,
) -> str:
    """
    Format a simple cost comparison table as a markdown string.

    Default costs:
      - OSS (self-hosted): $0.00 per 1k tokens
      - Llama 3.1-8B via Groq: $0.00 (free tier) / ~$0.00005 per 1k tokens (paid)
    """
    oss_cost = (oss_tokens / 1000) * oss_cost_per_1k
    frontier_cost = (frontier_tokens / 1000) * frontier_cost_per_1k

    return (
        f"| Model       | Tokens | Cost (USD) |\n"
        f"|-------------|--------|------------|\n"
        f"| OSS (local) | {oss_tokens:,} | ${oss_cost:.4f} |\n"
        f"| Frontier    | {frontier_tokens:,} | ${frontier_cost:.4f} |\n"
    )