"""
Cost Tracker
Tracks and estimates token costs across models.
"""
import time
from typing import Dict, Any, List
import sqlite3
from .usage_tracker import UsageTracker

MODEL_PRICING = {
    "claude-opus-4-20250514": {"input_per_1k": 0.015, "output_per_1k": 0.075},
    "claude-sonnet-4-20250514": {"input_per_1k": 0.003, "output_per_1k": 0.015},
    "claude-haiku-3-5-20241022": {"input_per_1k": 0.00025, "output_per_1k": 0.00125},
    "claude": {"input_per_1k": 0.003, "output_per_1k": 0.015}, # default to sonnet
    "gpt-4o": {"input_per_1k": 0.005, "output_per_1k": 0.015},
    "gpt-4o-mini": {"input_per_1k": 0.00015, "output_per_1k": 0.0006},
    "o3-mini": {"input_per_1k": 0.0011, "output_per_1k": 0.0044},
    "gemini-2.0-flash": {"input_per_1k": 0.0001, "output_per_1k": 0.0004},
    "gemini-2.5-pro": {"input_per_1k": 0.00125, "output_per_1k": 0.005},
    "gemini": {"input_per_1k": 0.00125, "output_per_1k": 0.005}, # default to pro
}

def estimate_cost(model: str, prompt_tokens: int, response_tokens: int) -> float:
    if model.startswith("ollama:"):
        return 0.0
    pricing = MODEL_PRICING.get(model, {"input_per_1k": 0.0, "output_per_1k": 0.0})
    return (prompt_tokens / 1000.0) * pricing["input_per_1k"] + (response_tokens / 1000.0) * pricing["output_per_1k"]

class CostTracker:
    def __init__(self, tracker: UsageTracker):
        self.tracker = tracker

    def daily_cost(self, days: int = 1) -> Dict[str, Any]:
        with sqlite3.connect(self.tracker.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT model, strategy, prompt_tokens, response_tokens
                         FROM interactions WHERE timestamp >= ?''', (cutoff,))
            
            total_usd = 0.0
            by_model = {}
            by_strategy = {}
            
            for row in c.fetchall():
                model, strat, p_tok, r_tok = row
                cost = estimate_cost(model, p_tok, r_tok)
                total_usd += cost
                by_model[model] = by_model.get(model, 0.0) + cost
                by_strategy[strat] = by_strategy.get(strat, 0.0) + cost
                
            return {
                "total_usd": total_usd,
                "by_model": by_model,
                "by_strategy": by_strategy
            }

    def weekly_report(self, weeks: int = 1) -> Dict[str, Any]:
        days = weeks * 7
        with sqlite3.connect(self.tracker.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT strftime('%Y-%m-%d', timestamp), model, prompt_tokens, response_tokens
                         FROM interactions WHERE timestamp >= ?''', (cutoff,))
            
            total_usd = 0.0
            daily = {}
            models = {}
            
            for row in c.fetchall():
                date, model, p_tok, r_tok = row
                cost = estimate_cost(model, p_tok, r_tok)
                total_usd += cost
                daily[date] = daily.get(date, 0.0) + cost
                models[model] = models.get(model, 0.0) + cost
                
            daily_breakdown = [{"date": k, "cost": v} for k, v in sorted(daily.items())]
            top_model = max(models.items(), key=lambda x: x[1])[0] if models else None
            cheap_model = min(models.items(), key=lambda x: x[1])[0] if models else None
            
            return {
                "total_usd": total_usd,
                "daily_breakdown": daily_breakdown,
                "top_cost_model": top_model,
                "cheapest_model": cheap_model,
                "avg_daily_cost": total_usd / float(days)
            }

    def monthly_report(self) -> Dict[str, Any]:
        # Approximate month as 30 days
        days = 30
        with sqlite3.connect(self.tracker.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT strftime('%W', timestamp), model, prompt_tokens, response_tokens
                         FROM interactions WHERE timestamp >= ?''', (cutoff,))
            
            total_usd = 0.0
            weekly = {}
            
            for row in c.fetchall():
                week, model, p_tok, r_tok = row
                cost = estimate_cost(model, p_tok, r_tok)
                total_usd += cost
                weekly[week] = weekly.get(week, 0.0) + cost
                
            weekly_breakdown = [{"week": k, "cost": v} for k, v in sorted(weekly.items())]
            
            # Get budget
            budget = self.check_budget()["budget"]
            budget_used_pct = (total_usd / (budget * 30)) * 100 if budget > 0 else None
            
            return {
                "total_usd": total_usd,
                "weekly_breakdown": weekly_breakdown,
                "budget_used_pct": budget_used_pct
            }

    def set_budget(self, daily_usd: float):
        from .preferences import Preferences
        prefs = Preferences(self.tracker.data_dir.parent)
        prefs.save("daily_budget_usd", daily_usd)

    def check_budget(self) -> Dict[str, Any]:
        from .preferences import Preferences
        prefs = Preferences(self.tracker.data_dir.parent)
        budget = prefs.get("daily_budget_usd", 0.0)
        spent = self.daily_cost(days=1)["total_usd"]
        return {
            "budget": budget,
            "spent_today": spent,
            "remaining": budget - spent,
            "over_budget": spent > budget if budget > 0 else False
        }

    def cost_per_chain(self, chain_result: Dict[str, Any]) -> Dict[str, Any]:
        total_usd = 0.0
        steps_cost = []
        for step in chain_result.get("steps", []):
            model = step.get("model", "unknown")
            tokens = step.get("tokens", 0) # Assuming this is sum of prompt+response
            # Approximate by halving since we don't have prompt/response split in result
            cost = estimate_cost(model, tokens // 2, tokens // 2)
            total_usd += cost
            steps_cost.append({"model": model, "tokens": tokens, "cost": cost})
        return {"total_usd": total_usd, "steps": steps_cost}

    def savings_from_ollama(self, days: int = 30) -> Dict[str, Any]:
        with sqlite3.connect(self.tracker.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT prompt_tokens, response_tokens
                         FROM interactions WHERE model LIKE 'ollama:%' AND timestamp >= ?''', (cutoff,))
            rows = c.fetchall()
            
            total_interactions = len(rows)
            estimated_cost = 0.0
            
            # Compare against claude-haiku as cheapest cloud
            for r in rows:
                estimated_cost += estimate_cost("claude-haiku-3-5-20241022", r[0], r[1])
                
            return {
                "interactions": total_interactions,
                "estimated_cloud_cost": estimated_cost,
                "actual_cost": 0.0,
                "saved_usd": estimated_cost
            }
