"""Usage Logging Utility"""
from sqlalchemy.orm import Session
from app.models.usage import UsageLog
from app.models.pricing import ModelPricing
from app.models.job import Job
from typing import Optional, Dict, Any
import logging

logger = logging.getLogger(__name__)

def log_model_usage(
    db: Session,
    job_id: str,
    user_id: str,
    module: str,
    action: str,
    provider: str,
    model: str,
    request_metadata: Dict[str, Any] = {},
    response_metadata: Dict[str, Any] = {},
    usage_stats: Dict[str, Any] = {}
) -> Optional[UsageLog]:
    """
    Log model usage and calculate estimated cost.
    
    usage_stats keys:
    - input_tokens (int)
    - output_tokens (int)
    - images (int)
    - seconds (float)
    - characters (int)
    """
    try:
        # 1. Calculate Cost
        cost = 0.0
        
        # Find pricing record
        pricing = db.query(ModelPricing).filter(
            ModelPricing.provider == provider,
            ModelPricing.model_name == model
        ).first()
        
        # If specific model not found, try generic provider/default fallback?
        # For now, just log 0 if not found, or maybe try mapping aliases.
        
        if pricing:
            if usage_stats.get("input_tokens"):
                cost += float(pricing.cost_per_input_token) * usage_stats["input_tokens"]
            
            if usage_stats.get("output_tokens"):
                cost += float(pricing.cost_per_output_token) * usage_stats["output_tokens"]
                
            if usage_stats.get("images"):
                cost += float(pricing.cost_per_image) * usage_stats["images"]
                
            if usage_stats.get("seconds"):
                cost += float(pricing.cost_per_second) * usage_stats["seconds"]
                
            if usage_stats.get("characters"):
                cost += float(pricing.cost_per_1k_chars) * (usage_stats["characters"] / 1000.0)
                
            if usage_stats.get("requests"): # generic per-request
                 cost += float(pricing.cost_per_request) * usage_stats["requests"]
            
            # Special case for "per request" if not specified but implies 1?
            # If generated 1 image and pricing is per image... handled above.
        
        # 2. Create Log Record
        log_entry = UsageLog(
            job_id=job_id,
            user_id=user_id,
            module=module,
            action=action,
            api_provider=provider,
            api_model=model,
            tokens_input=usage_stats.get("input_tokens"),
            tokens_output=usage_stats.get("output_tokens"),
            estimated_cost_usd=cost,
            processing_time_ms=usage_stats.get("processing_time_ms"),
            request_metadata=request_metadata,
            response_metadata=response_metadata
        )
        
        db.add(log_entry)
        db.commit()
        return log_entry
        
    except Exception as e:
        logger.error(f"Failed to log usage: {e}")
        return None