import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

COST_LOG_PATH = Path("data/cost_log.jsonl")

# Pricing per token, by model id. Anthropic prices per platform.claude.com/docs/pricing;
# OpenAI prices approximate — verify against platform.openai.com/pricing before relying
# on this for exact billing reconciliation.
_PRICING = {
    "claude-opus-4-8":            {"input": 5.0 / 1_000_000,  "output": 25.0 / 1_000_000},
    "claude-sonnet-4-6":          {"input": 3.0 / 1_000_000,  "output": 15.0 / 1_000_000},
    "claude-haiku-4-5-20251001":  {"input": 1.0 / 1_000_000,  "output": 5.0 / 1_000_000},
    "gpt-5.1":                    {"input": 1.25 / 1_000_000, "output": 10.0 / 1_000_000},
    "gpt-5-mini":                 {"input": 0.25 / 1_000_000, "output": 2.0 / 1_000_000},
}
_DEFAULT_PRICING = {"input": 3.0 / 1_000_000, "output": 15.0 / 1_000_000}


def log_usage(ticker: str, model: str, input_tokens: int, output_tokens: int) -> None:
    """Append one usage record to data/cost_log.jsonl."""
    COST_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
    pricing = _PRICING.get(model, _DEFAULT_PRICING)
    cost_usd = input_tokens * pricing["input"] + output_tokens * pricing["output"]
    record = {
        "ts": datetime.now(timezone.utc).isoformat(),
        "ticker": ticker,
        "model": model,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "cost_usd": round(cost_usd, 6),
    }
    with COST_LOG_PATH.open("a", encoding="utf-8") as f:
        f.write(json.dumps(record) + "\n")


def _message_usage(msg) -> tuple[int, int]:
    """Extract (input_tokens, output_tokens) from one AI message.

    Tries three shapes, in order, and stops at the first that yields tokens
    (never double-counts a message across shapes):
      1. `usage_metadata` — LangChain's provider-normalized field, present on
         both langchain-anthropic and langchain-openai AIMessages.
      2. Anthropic legacy: `response_metadata['usage']`.
      3. OpenAI legacy: `response_metadata['token_usage']` (prompt/completion).
    """
    usage_metadata = getattr(msg, "usage_metadata", None) or {}
    if usage_metadata.get("input_tokens") or usage_metadata.get("output_tokens"):
        return usage_metadata.get("input_tokens", 0), usage_metadata.get("output_tokens", 0)

    meta = getattr(msg, "response_metadata", {}) or {}

    anthropic_usage = meta.get("usage") or {}
    if anthropic_usage.get("input_tokens") or anthropic_usage.get("output_tokens"):
        return anthropic_usage.get("input_tokens", 0), anthropic_usage.get("output_tokens", 0)

    openai_usage = meta.get("token_usage") or {}
    if openai_usage.get("prompt_tokens") or openai_usage.get("completion_tokens"):
        return openai_usage.get("prompt_tokens", 0), openai_usage.get("completion_tokens", 0)

    return 0, 0


def compute_run_cost(messages: list, model: str) -> dict:
    """Extract usage from LangChain AI messages and return total cost summary."""
    total_input = 0
    total_output = 0
    for msg in messages:
        in_tok, out_tok = _message_usage(msg)
        total_input += in_tok
        total_output += out_tok

    pricing = _PRICING.get(model, _DEFAULT_PRICING)
    cost_usd = total_input * pricing["input"] + total_output * pricing["output"]
    return {
        "input_tokens": total_input,
        "output_tokens": total_output,
        "cost_usd": round(cost_usd, 6),
    }


def average_brief_cost_usd() -> Optional[float]:
    """Return average cost_usd per brief run from cost_log.jsonl, or None if log is empty.

    Used by the Portfolio screen to show an estimated cost before bulk generation.
    """
    if not COST_LOG_PATH.exists():
        return None
    costs: list[float] = []
    try:
        with COST_LOG_PATH.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                record = json.loads(line)
                c = record.get("cost_usd")
                if c is not None:
                    costs.append(float(c))
    except Exception:
        return None
    return sum(costs) / len(costs) if costs else None