"""Append-only log of AI-provider usage (tokens + Tavily credits) per site.

The cron generators (articles, comments, fact-checks, corrections) and their
Tavily research all record usage here; the /usage endpoint reads it back and
aggregates per model / per provider with estimated cost.

Storage mirrors `posting_history.py`: a per-environment JSONL file under
`app/data/` — one JSON object per line — so independent processes (the API and
the hourly cron timers) can append without a read-modify-write race. Recording
is best-effort and never raises into the caller: instrumentation must never
break generation.

Attribution uses an ambient site slug set by a contextvar at each flow boundary
(`ArticleCreationService.create_article`, the comment endpoint / agent,
`run_corrections_core`). This keeps the record calls inside `ChatService` down
to one line each without threading `site` through every method signature. The
contextvar propagates to code run inline and via `asyncio.to_thread` /
Starlette's threadpool, which covers every current call path.
"""

import contextvars
import csv
import io
import json
import logging
import os
from contextlib import contextmanager
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Iterator, List, Optional

from ..appTypes import (
    DailyUsage,
    ModelUsage,
    OperationUsage,
    ProviderUsage,
    UsageDetailResponse,
    UsageRecord,
    UsageSummaryResponse,
)
from ..data import pricing

CACHE_DIR = Path(__file__).parent.parent / 'data'

# Ambient site slug for the current generation flow. Set at boundaries via
# usage_scope(); read by record_usage() when no explicit slug is passed.
_current_site_slug: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    'tsai_usage_site_slug', default=None
)

# Provider display names (match LLMProvider enum values + Tavily).
PROVIDER_OPENAI = 'OpenAI'
PROVIDER_CLAUDE = 'Claude'
PROVIDER_TAVILY = 'Tavily'


def _usage_path() -> Path:
    """Env-suffixed usage file, matching SiteConfigService's cache mapping."""
    env = os.getenv('TSAI_ENV', 'development')
    if env == 'production':
        suffix = 'tsai'
    elif env == 'scike_ai':
        suffix = 'scike'
    else:
        suffix = 'development'
    return CACHE_DIR / f'usage_{suffix}.jsonl'


@contextmanager
def usage_scope(site_slug: Optional[str]) -> Iterator[None]:
    """Set the ambient site slug for the duration of a generation flow.

    Nesting is fine — inner scopes reset to the outer value on exit.
    """
    token = _current_site_slug.set(site_slug)
    try:
        yield
    finally:
        _current_site_slug.reset(token)


def enter_usage_scope(site_slug: Optional[str]):
    """Set the ambient slug, returning a token to pass to exit_usage_scope().

    Use this to instrument a method that already has its own try/finally without
    re-indenting the body under `with usage_scope(...)`.
    """
    return _current_site_slug.set(site_slug)


def exit_usage_scope(token) -> None:
    try:
        _current_site_slug.reset(token)
    except Exception:
        pass


def _ambient_slug() -> Optional[str]:
    return _current_site_slug.get()


def record_usage(
    *,
    provider: str,
    model: str,
    operation: str,
    input_tokens: int = 0,
    output_tokens: int = 0,
    tavily_credits: int = 0,
    images: int = 0,
    image_size: Optional[str] = None,
    image_quality: Optional[str] = None,
    site_slug: Optional[str] = None,
    when: Optional[datetime] = None,
) -> None:
    """Append one usage event. Best-effort — never raises.

    Falls back to the ambient slug when site_slug is not given; skips silently
    when there is no slug to attribute to (e.g. a call outside any usage_scope).
    """
    try:
        slug = site_slug or _ambient_slug()
        if not slug:
            return
        if not (input_tokens or output_tokens or tavily_credits or images):
            return
        now = when or datetime.now()
        record = UsageRecord(
            ts=now.isoformat(timespec='seconds'),
            date=now.strftime('%Y-%m-%d'),
            site=slug,
            provider=provider,
            model=model or 'unknown',
            operation=operation,
            input_tokens=int(input_tokens or 0),
            output_tokens=int(output_tokens or 0),
            tavily_credits=int(tavily_credits or 0),
            images=int(images or 0),
            image_size=image_size,
            image_quality=image_quality,
        )
        path = _usage_path()
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'a', encoding='utf-8') as f:
            f.write(json.dumps(record.model_dump()) + '\n')
            f.flush()
    except Exception as e:
        logging.error(f"Failed to record usage ({provider}/{model}/{operation}): {e}")


def _resolve_image_size(size: Optional[str], b64: Optional[str]) -> Optional[str]:
    """Resolve an image size to a concrete "WxH".

    Concrete sizes pass through. 'auto'/None are measured from the returned
    base64 image via PIL; if that fails, we return 'auto' (priced as unknown).
    """
    s = (size or '').lower()
    if s and s != 'auto' and 'x' in s:
        return size
    if not b64:
        return size or 'auto'
    try:
        import base64 as _b64
        import io
        from PIL import Image
        raw = _b64.b64decode(b64)
        with Image.open(io.BytesIO(raw)) as im:
            w, h = im.size
        return f"{w}x{h}"
    except Exception as e:
        logging.warning(f"Could not measure image size (falling back to '{size}'): {e}")
        return size or 'auto'


def record_image(
    *,
    model: str,
    size: Optional[str],
    quality: Optional[str],
    operation: str,
    provider: str = PROVIDER_OPENAI,
    count: int = 1,
    b64: Optional[str] = None,
    site_slug: Optional[str] = None,
) -> None:
    """Record one (or `count`) generated image(s). Best-effort — never raises.

    Resolves an 'auto' size by measuring `b64` when provided so the per-image
    price table can be applied.
    """
    resolved_size = _resolve_image_size(size, b64)
    record_usage(
        provider=provider,
        model=model,
        operation=operation,
        images=count,
        image_size=resolved_size,
        image_quality=(quality or None),
        site_slug=site_slug,
    )


def record_openai_completion(completion: Any, *, model: str, operation: str) -> None:
    """Record token usage from an OpenAI completion (chat.completions / .parse)."""
    try:
        usage = getattr(completion, 'usage', None)
        if not usage:
            return
        record_usage(
            provider=PROVIDER_OPENAI,
            model=model,
            operation=operation,
            input_tokens=getattr(usage, 'prompt_tokens', 0) or 0,
            output_tokens=getattr(usage, 'completion_tokens', 0) or 0,
        )
    except Exception as e:
        logging.error(f"Failed to record OpenAI usage ({model}/{operation}): {e}")


def record_langchain_message(message: Any, *, provider: str, model: str, operation: str, site_slug: Optional[str] = None) -> None:
    """Record token usage from a LangChain AIMessage/AIMessageChunk (.usage_metadata).

    Pass site_slug explicitly for streaming call sites that aren't inside a
    usage_scope; otherwise the ambient slug is used.
    """
    try:
        meta = getattr(message, 'usage_metadata', None) or {}
        input_tokens = meta.get('input_tokens', 0) or 0
        output_tokens = meta.get('output_tokens', 0) or 0
        record_usage(
            provider=provider,
            model=model,
            operation=operation,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            site_slug=site_slug,
        )
    except Exception as e:
        logging.error(f"Failed to record LangChain usage ({provider}/{model}/{operation}): {e}")


def record_agent_usage(
    callback: Any, *, provider: str, operation: str, site_slug: Optional[str] = None
) -> None:
    """Record token usage collected by a LangChain UsageMetadataCallbackHandler.

    An AgentExecutor runs its driver model across several tool-calling iterations;
    the handler aggregates each turn's usage_metadata into a dict keyed by model
    name (``{model: {input_tokens, output_tokens, ...}}``). We record one event per
    model. Best-effort — never raises into the caller.
    """
    try:
        for model_name, meta in (getattr(callback, 'usage_metadata', {}) or {}).items():
            meta = meta or {}
            record_usage(
                provider=provider,
                model=model_name,
                operation=operation,
                input_tokens=meta.get('input_tokens', 0) or 0,
                output_tokens=meta.get('output_tokens', 0) or 0,
                site_slug=site_slug,
            )
    except Exception as e:
        logging.error(f"Failed to record agent usage ({provider}/{operation}): {e}")


def record_tavily_search(*, search_depth: str = 'basic', operation: str = 'research') -> None:
    """Record one Tavily search as credits (basic = 1, advanced = 2)."""
    credits = 2 if (search_depth or 'basic').lower() == 'advanced' else 1
    record_usage(
        provider=PROVIDER_TAVILY,
        model='search',
        operation=operation,
        tavily_credits=credits,
    )


def used_credits_since(site_slug: str, since_ts: datetime) -> float:
    """Estimated credits (1 credit = $1) this site has used since `since_ts`.

    Powers credit enforcement: the store keeps the billing period boundary, the
    API converts the append-only usage log into a spend level since that
    boundary. Reads a slightly padded day window then filters precisely by
    timestamp so partial first/last days aren't truncated.
    """
    now = datetime.now()
    span_days = max(1, (now - since_ts).days + 2)
    records = read_usage(site_slug, since_days=span_days)
    cutoff_iso = since_ts.isoformat(timespec='seconds')
    in_period = [r for r in records if (r.ts or '') >= cutoff_iso]
    return _price_records(in_period)


def read_usage(site_slug: str, *, since_days: int = 30) -> List[UsageRecord]:
    """Return this site's usage records within the last `since_days` days.

    Tolerant of malformed / partially-written trailing lines — those are skipped.
    """
    path = _usage_path()
    if not path.exists():
        return []

    cutoff = (datetime.now() - timedelta(days=since_days)).strftime('%Y-%m-%d')
    records: List[UsageRecord] = []
    try:
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    data = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if data.get('site') != site_slug:
                    continue
                if str(data.get('date', '')) < cutoff:
                    continue
                try:
                    records.append(UsageRecord(**data))
                except Exception:
                    continue
    except Exception as e:
        logging.error(f"Failed to read usage for {site_slug}: {e}")
        return records

    return records


# Provider display order for the card (paid consumers first, Tavily last).
_PROVIDER_ORDER = {PROVIDER_OPENAI: 0, PROVIDER_CLAUDE: 1, PROVIDER_TAVILY: 2}


def _bucket_records(records: List[UsageRecord]) -> dict:
    """Group records by pricing key: (provider, model) for tokens/credits, or
    (provider, model, 'img', size, quality) for images (priced per image)."""
    buckets: dict = {}
    for r in records:
        if r.images:
            key = (r.provider, r.model, 'img', r.image_size, r.image_quality)
        else:
            key = (r.provider, r.model)
        acc = buckets.setdefault(
            key, {'input': 0, 'output': 0, 'credits': 0, 'images': 0}
        )
        acc['input'] += r.input_tokens
        acc['output'] += r.output_tokens
        acc['credits'] += r.tavily_credits
        acc['images'] += r.images
    return buckets


def _price_bucket(key: tuple, acc: dict) -> tuple:
    """Price one bucket produced by _bucket_records. Returns (cost, price_known)."""
    if len(key) == 5:  # image key
        provider, model, _, image_size, image_quality = key
    else:
        provider, model = key
        image_size = image_quality = None
    return pricing.estimate_cost(
        provider,
        model,
        input_tokens=acc['input'],
        output_tokens=acc['output'],
        tavily_credits=acc['credits'],
        images=acc['images'],
        image_size=image_size,
        image_quality=image_quality,
    )


def _price_records(records: List[UsageRecord]) -> float:
    """Total estimated USD cost for a set of records, using the same per-model /
    per-image bucketing as aggregate_usage so day/operation costs reconcile with
    the provider totals."""
    total = 0.0
    for key, acc in _bucket_records(records).items():
        cost, _ = _price_bucket(key, acc)
        total += cost
    return round(total, 6)


def aggregate_usage(
    records: List[UsageRecord], slug: str, since_days: int
) -> UsageSummaryResponse:
    """Group records into rows, price each, and roll up per-provider subtotals.

    Token/Tavily rows group by (provider, model). Image rows group by
    (provider, model, size, quality) and are priced per image — the size/quality
    are folded into a display label so distinct tiers show as their own rows.
    """
    # provider -> list[ModelUsage]
    by_provider: dict = {}
    total_cost = 0.0
    for key, acc in _bucket_records(records).items():
        is_image = len(key) == 5
        if is_image:
            provider, model, _, image_size, image_quality = key
            label = f"{model} ({image_size or 'auto'}, {image_quality or '?'})"
        else:
            provider, model = key
            label = model

        cost, price_known = _price_bucket(key, acc)
        total_cost += cost
        by_provider.setdefault(provider, []).append(
            ModelUsage(
                provider=provider,
                model=label,
                input_tokens=acc['input'],
                output_tokens=acc['output'],
                tavily_credits=acc['credits'],
                images=acc['images'],
                estimated_cost=cost,
                price_known=price_known,
            )
        )

    providers: List[ProviderUsage] = []
    for provider in sorted(by_provider, key=lambda p: _PROVIDER_ORDER.get(p, 99)):
        models = sorted(by_provider[provider], key=lambda m: m.estimated_cost, reverse=True)
        subtotal = round(sum(m.estimated_cost for m in models), 6)
        providers.append(
            ProviderUsage(provider=provider, models=models, subtotal_cost=subtotal)
        )

    return UsageSummaryResponse(
        slug=slug,
        since_days=since_days,
        providers=providers,
        total_cost=round(total_cost, 6),
        currency='USD',
        generated_at=datetime.now().isoformat(timespec='seconds'),
    )


def aggregate_daily(records: List[UsageRecord]) -> List[DailyUsage]:
    """Estimated spend per calendar day, ascending by date. Only days with
    records appear (the bar view scales itself; no gap-filling needed)."""
    by_date: dict = {}
    for r in records:
        by_date.setdefault(r.date, []).append(r)
    return [
        DailyUsage(date=date, cost=_price_records(day_records))
        for date, day_records in sorted(by_date.items())
    ]


def aggregate_by_operation(records: List[UsageRecord]) -> List[OperationUsage]:
    """Estimated spend rolled up by operation type, sorted by cost descending."""
    by_op: dict = {}
    for r in records:
        acc = by_op.setdefault(
            r.operation or 'unknown',
            {'records': [], 'input': 0, 'output': 0, 'images': 0, 'credits': 0},
        )
        acc['records'].append(r)
        acc['input'] += r.input_tokens
        acc['output'] += r.output_tokens
        acc['images'] += r.images
        acc['credits'] += r.tavily_credits

    rows = [
        OperationUsage(
            operation=op,
            cost=_price_records(acc['records']),
            input_tokens=acc['input'],
            output_tokens=acc['output'],
            images=acc['images'],
            tavily_credits=acc['credits'],
        )
        for op, acc in by_op.items()
    ]
    return sorted(rows, key=lambda o: o.cost, reverse=True)


def aggregate_detail(
    records: List[UsageRecord], slug: str, since_days: int
) -> UsageDetailResponse:
    """Full detail: per-provider/model rows plus daily and per-operation breakdowns.

    Reuses aggregate_usage for the provider rows and total so the detail page and
    the home card stay perfectly consistent."""
    summary = aggregate_usage(records, slug, since_days)
    return UsageDetailResponse(
        slug=slug,
        since_days=since_days,
        providers=summary.providers,
        daily=aggregate_daily(records),
        operations=aggregate_by_operation(records),
        total_cost=summary.total_cost,
        currency='USD',
        generated_at=datetime.now().isoformat(timespec='seconds'),
    )


# Columns for the raw per-event CSV export (billing / reconciliation).
_CSV_FIELDS = [
    'ts', 'date', 'site', 'provider', 'model', 'operation',
    'input_tokens', 'output_tokens', 'tavily_credits',
    'images', 'image_size', 'image_quality', 'estimated_cost',
]


def records_to_csv(records: List[UsageRecord]) -> str:
    """Serialize raw usage events to CSV, one row per recorded event, ascending by
    timestamp. Each row carries a per-event estimated cost (priced the same way as
    the aggregated views) so an external system can reconcile spend line-by-line."""
    out = io.StringIO()
    writer = csv.DictWriter(out, fieldnames=_CSV_FIELDS)
    writer.writeheader()
    for r in sorted(records, key=lambda x: x.ts):
        cost, _ = pricing.estimate_cost(
            r.provider,
            r.model,
            input_tokens=r.input_tokens,
            output_tokens=r.output_tokens,
            tavily_credits=r.tavily_credits,
            images=r.images,
            image_size=r.image_size,
            image_quality=r.image_quality,
        )
        writer.writerow({
            'ts': r.ts,
            'date': r.date,
            'site': r.site,
            'provider': r.provider,
            'model': r.model,
            'operation': r.operation,
            'input_tokens': r.input_tokens,
            'output_tokens': r.output_tokens,
            'tavily_credits': r.tavily_credits,
            'images': r.images,
            'image_size': r.image_size or '',
            'image_quality': r.image_quality or '',
            'estimated_cost': cost,
        })
    return out.getvalue()