#!/usr/bin/env python3
from __future__ import annotations

import json
import os
import re
import sys
from collections import defaultdict
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

import httpx
from dotenv import load_dotenv


ROOT = Path(r"C:\Users\Paul\OneDrive\Documents\IterFact")
MCP_DIR = ROOT / "MCP"
ROOT_ENV = ROOT / ".env.supabase.local"
MCP_ENV = MCP_DIR / ".env"

if ROOT_ENV.exists():
    load_dotenv(ROOT_ENV)
if MCP_ENV.exists():
    load_dotenv(MCP_ENV, override=False)


def env(name: str, default: str = "") -> str:
    return os.getenv(name, default).strip()


def require(name: str) -> str:
    value = env(name)
    if not value or value.startswith("your-"):
        raise RuntimeError(f"Missing required environment value: {name}")
    return value


def now_utc() -> datetime:
    return datetime.now(UTC)


def iso_day(value: datetime) -> str:
    return value.strftime("%Y-%m-%d")


def as_number(value: Any) -> float:
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        try:
            return float(value)
        except ValueError:
            return 0.0
    return 0.0


@dataclass
class Settings:
    anthropic_admin_api_key: str = require("ANTHROPIC_ADMIN_API_KEY")
    anthropic_version: str = env("ANTHROPIC_VERSION", "2023-06-01")
    supabase_url: str = require("NEXT_PUBLIC_SUPABASE_URL").rstrip("/")
    supabase_service_key: str = env("SUPABASE_SECRET_KEY") or require(
        "SUPABASE_SERVICE_ROLE_KEY"
    )


SETTINGS = Settings()


class SupabaseRest:
    def __init__(self, url: str, service_key: str) -> None:
        self.base_url = f"{url}/rest/v1"
        self.headers = {
            "apikey": service_key,
            "Authorization": f"Bearer {service_key}",
            "Content-Type": "application/json",
        }

    def upsert(self, table: str, rows: list[dict[str, Any]], conflict_column: str) -> None:
        if not rows:
            return

        response = httpx.post(
            f"{self.base_url}/{table}",
            headers={
                **self.headers,
                "Prefer": "resolution=merge-duplicates,return=minimal",
            },
            params={"on_conflict": conflict_column},
            json=rows,
            timeout=120,
        )
        response.raise_for_status()

    def select(
        self,
        table: str,
        select_clause: str,
        params: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        response = httpx.get(
            f"{self.base_url}/{table}",
            headers=self.headers,
            params={"select": select_clause, **(params or {})},
            timeout=120,
        )
        response.raise_for_status()
        return response.json()

    def update_by_id(self, table: str, row_id: str, payload: dict[str, Any]) -> None:
        response = httpx.patch(
            f"{self.base_url}/{table}",
            headers={**self.headers, "Prefer": "return=minimal"},
            params={"id": f"eq.{row_id}"},
            json=payload,
            timeout=120,
        )
        response.raise_for_status()


class AnthropicAdmin:
    def __init__(self, admin_key: str, version: str) -> None:
        self.headers = {
            "x-api-key": admin_key,
            "anthropic-version": version,
        }
        self.base_url = "https://api.anthropic.com/v1/organizations"

    def fetch_report(
        self,
        endpoint: str,
        starting_at: str,
        ending_at: str,
        bucket_width: str,
        group_by: list[str],
    ) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        next_page: str | None = None

        while True:
            params: list[tuple[str, str]] = [
                ("starting_at", starting_at),
                ("ending_at", ending_at),
                ("bucket_width", bucket_width),
            ]
            params.extend(("group_by[]", item) for item in group_by)
            if next_page:
                params.append(("page", next_page))

            response = httpx.get(
                f"{self.base_url}/{endpoint}",
                headers=self.headers,
                params=params,
                timeout=120,
            )
            response.raise_for_status()
            payload = response.json()
            results.extend(flatten_buckets(payload))

            next_page = (
                payload.get("next_page")
                or payload.get("next_cursor")
                or payload.get("next_token")
            )
            if not next_page and not payload.get("has_more"):
                break
            if payload.get("has_more") and not next_page:
                break

        return results


def flatten_buckets(payload: dict[str, Any]) -> list[dict[str, Any]]:
    raw_items = payload.get("data") or payload.get("results") or payload.get("items") or []
    flattened: list[dict[str, Any]] = []

    for raw in raw_items:
        if not isinstance(raw, dict):
            continue
        nested = raw.get("results")
        if isinstance(nested, list):
            bucket_meta = {
                key: raw.get(key)
                for key in (
                    "bucket_width",
                    "starting_at",
                    "ending_at",
                    "start_time",
                    "end_time",
                    "bucket_date",
                    "currency",
                )
            }
            for entry in nested:
                if isinstance(entry, dict):
                    flattened.append({**bucket_meta, **entry, "raw_bucket": raw})
            continue

        flattened.append(raw)

    return flattened


def normalize_bucket_date(raw: dict[str, Any]) -> str:
    for key in ("bucket_date", "starting_at", "start_time", "date"):
        value = raw.get(key)
        if isinstance(value, str) and value:
            return value[:10]
    return iso_day(now_utc())


def normalize_group_value(value: Any) -> str | None:
    if value is None:
        return None
    if isinstance(value, dict):
        for key in ("id", "name", "value", "workspace_id", "api_key", "model", "service_tier"):
            if key in value and value[key]:
                return str(value[key])
        return json.dumps(value, sort_keys=True)
    if isinstance(value, (list, tuple)):
        joined = ",".join(str(item) for item in value if item is not None)
        return joined or None
    text = str(value).strip()
    return text or None


def extract_model(raw: dict[str, Any]) -> str | None:
    direct = normalize_group_value(raw.get("model"))
    if direct:
        return direct

    description = raw.get("description")
    if isinstance(description, dict):
        for key in ("model", "display_name", "name"):
            value = normalize_group_value(description.get(key))
            if value:
                return value
    elif isinstance(description, str) and "claude" in description.lower():
        return description

    return None


def extract_service_tier(raw: dict[str, Any]) -> str | None:
    direct = normalize_group_value(raw.get("service_tier"))
    if direct:
        return direct

    description = raw.get("description")
    if isinstance(description, dict):
        return normalize_group_value(description.get("service_tier"))

    return None


def extract_workspace_id(raw: dict[str, Any]) -> str | None:
    return normalize_group_value(raw.get("workspace_id") or raw.get("workspace"))


def extract_api_key_label(raw: dict[str, Any]) -> str | None:
    for key in ("api_key", "api_key_id", "api_key_name"):
        value = normalize_group_value(raw.get(key))
        if value:
            return value
    return None


def extract_usage_metrics(raw: dict[str, Any]) -> dict[str, float]:
    return {
        "input_tokens": as_number(raw.get("input_tokens") or raw.get("total_input_tokens")),
        "output_tokens": as_number(raw.get("output_tokens") or raw.get("total_output_tokens")),
        "cache_creation_input_tokens": as_number(
            raw.get("cache_creation_input_tokens") or raw.get("cache_creation_tokens")
        ),
        "cache_read_input_tokens": as_number(
            raw.get("cache_read_input_tokens") or raw.get("cache_read_tokens")
        ),
        "request_count": as_number(raw.get("request_count") or raw.get("count")),
    }


def extract_cost_metrics(raw: dict[str, Any]) -> tuple[float, str]:
    amount = raw.get("amount") or raw.get("cost") or raw.get("total_cost")
    currency = normalize_group_value(raw.get("currency")) or "USD"

    if isinstance(amount, dict):
        value = (
            amount.get("usd")
            or amount.get("value")
            or amount.get("amount")
            or amount.get("cost")
        )
        currency = normalize_group_value(amount.get("currency")) or currency
        return as_number(value), currency

    if amount is not None:
        return as_number(amount), currency

    for key in ("cost_usd", "total_cost_usd", "amount_usd"):
        if raw.get(key) is not None:
            return as_number(raw.get(key)), "USD"

    return 0.0, currency


def build_usage_bucket_row(raw: dict[str, Any]) -> dict[str, Any]:
    metrics = extract_usage_metrics(raw)
    bucket_date = normalize_bucket_date(raw)
    model = extract_model(raw)
    workspace_id = extract_workspace_id(raw)
    service_tier = extract_service_tier(raw)
    api_key_label = extract_api_key_label(raw)

    bucket_key = "|".join(
        [
            "anthropic",
            "usage_report_messages",
            bucket_date,
            model or "any-model",
            workspace_id or "any-workspace",
            service_tier or "any-tier",
            api_key_label or "any-key",
        ]
    )

    return {
        "bucket_key": bucket_key,
        "provider": "anthropic",
        "source_endpoint": "usage_report_messages",
        "bucket_date": bucket_date,
        "bucket_width": normalize_group_value(raw.get("bucket_width")) or "1d",
        "model": model,
        "api_key_label": api_key_label,
        "workspace_id": workspace_id,
        "service_tier": service_tier,
        "total_input_tokens": int(metrics["input_tokens"]),
        "total_output_tokens": int(metrics["output_tokens"]),
        "total_cache_creation_input_tokens": int(metrics["cache_creation_input_tokens"]),
        "total_cache_read_input_tokens": int(metrics["cache_read_input_tokens"]),
        "request_count": int(metrics["request_count"]),
        "raw_bucket": raw,
        "synced_at": now_utc().isoformat(),
    }


def build_cost_bucket_row(raw: dict[str, Any]) -> dict[str, Any]:
    total_cost_usd, currency = extract_cost_metrics(raw)
    bucket_date = normalize_bucket_date(raw)
    model = extract_model(raw)
    workspace_id = extract_workspace_id(raw)
    service_tier = extract_service_tier(raw)
    group_values = {
        "description": raw.get("description"),
        "workspace_id": workspace_id,
        "model": model,
        "service_tier": service_tier,
    }

    bucket_key = "|".join(
        [
            "anthropic",
            "cost_report",
            bucket_date,
            model or "any-model",
            workspace_id or "any-workspace",
            service_tier or "any-tier",
            normalize_group_value(raw.get("description")) or "any-description",
        ]
    )

    return {
        "bucket_key": bucket_key,
        "provider": "anthropic",
        "source_endpoint": "cost_report",
        "bucket_date": bucket_date,
        "bucket_width": normalize_group_value(raw.get("bucket_width")) or "1d",
        "billing_currency": currency,
        "total_cost_usd": total_cost_usd,
        "workspace_id": workspace_id,
        "group_values": group_values,
        "raw_bucket": raw,
        "synced_at": now_utc().isoformat(),
    }


def event_match_keys(event: dict[str, Any]) -> list[tuple[str, str | None, str | None]]:
    bucket_date = normalize_bucket_date(event)
    model = normalize_group_value(event.get("model"))
    service_tier = normalize_group_value(event.get("service_tier"))
    return [
        (bucket_date, model, service_tier),
        (bucket_date, model, None),
        (bucket_date, None, service_tier),
        (bucket_date, None, None),
    ]


def bucket_match_keys(bucket: dict[str, Any]) -> list[tuple[str, str | None, str | None]]:
    bucket_date = bucket.get("bucket_date")
    group_values = bucket.get("group_values") or {}
    if not isinstance(group_values, dict):
        group_values = {}
    model = normalize_group_value(group_values.get("model"))
    service_tier = normalize_group_value(group_values.get("service_tier"))
    return [
        (bucket_date, model, service_tier),
        (bucket_date, model, None),
        (bucket_date, None, service_tier),
        (bucket_date, None, None),
    ]


def allocate_actual_costs(
    events: list[dict[str, Any]],
    buckets: list[dict[str, Any]],
) -> tuple[dict[str, float], dict[str, float]]:
    bucket_totals: dict[tuple[str, str | None, str | None], float] = defaultdict(float)
    for bucket in buckets:
        for key in bucket_match_keys(bucket):
            bucket_totals[key] += as_number(bucket.get("total_cost_usd"))
            break

    grouped_events: dict[tuple[str, str | None, str | None], list[dict[str, Any]]] = defaultdict(list)
    for event in events:
        for key in event_match_keys(event):
            if key in bucket_totals:
                grouped_events[key].append(event)
                break

    reconciled_by_event: dict[str, float] = {}
    reconciled_by_run: dict[str, float] = defaultdict(float)

    for key, group in grouped_events.items():
        actual_total = bucket_totals.get(key, 0.0)
        if actual_total <= 0:
            continue

        formula_total = sum(as_number(item.get("formula_cost_usd")) for item in group)
        token_total = sum(
            as_number(item.get("input_tokens"))
            + as_number(item.get("output_tokens"))
            + as_number(item.get("cache_creation_input_tokens"))
            + as_number(item.get("cache_read_input_tokens"))
            for item in group
        )

        running_total = 0.0
        for index, event in enumerate(group):
            denominator = formula_total if formula_total > 0 else token_total
            numerator = (
                as_number(event.get("formula_cost_usd"))
                if formula_total > 0
                else as_number(event.get("input_tokens"))
                + as_number(event.get("output_tokens"))
                + as_number(event.get("cache_creation_input_tokens"))
                + as_number(event.get("cache_read_input_tokens"))
            )

            if denominator <= 0:
                allocated = actual_total / len(group)
            else:
                allocated = actual_total * (numerator / denominator)

            if index == len(group) - 1:
                allocated = max(actual_total - running_total, 0.0)

            running_total += allocated
            event_id = event["id"]
            run_id = event["generation_run_id"]
            reconciled_by_event[event_id] = allocated
            reconciled_by_run[run_id] += allocated

    return reconciled_by_event, reconciled_by_run


def main() -> int:
    days_back = int(sys.argv[1]) if len(sys.argv) > 1 else 7
    ending_at = now_utc()
    starting_at = ending_at - timedelta(days=days_back)

    supabase = SupabaseRest(SETTINGS.supabase_url, SETTINGS.supabase_service_key)
    anthropic = AnthropicAdmin(
        SETTINGS.anthropic_admin_api_key,
        SETTINGS.anthropic_version,
    )

    usage_rows = anthropic.fetch_report(
        "usage_report/messages",
        starting_at.isoformat().replace("+00:00", "Z"),
        ending_at.isoformat().replace("+00:00", "Z"),
        "1d",
        ["model", "api_key", "workspace_id", "service_tier"],
    )
    cost_rows = anthropic.fetch_report(
        "cost_report",
        starting_at.isoformat().replace("+00:00", "Z"),
        ending_at.isoformat().replace("+00:00", "Z"),
        "1d",
        ["description", "workspace_id"],
    )

    usage_bucket_rows = [build_usage_bucket_row(row) for row in usage_rows]
    cost_bucket_rows = [build_cost_bucket_row(row) for row in cost_rows]
    supabase.upsert("provider_usage_buckets", usage_bucket_rows, "bucket_key")
    supabase.upsert("provider_cost_buckets", cost_bucket_rows, "bucket_key")

    event_rows = supabase.select(
        "provider_usage_events",
        "id,generation_run_id,model,service_tier,formula_cost_usd,input_tokens,output_tokens,cache_creation_input_tokens,cache_read_input_tokens,usage_bucket_date,reconciliation_status",
        {
            "provider": "eq.anthropic",
            "usage_bucket_date": f"gte.{iso_day(starting_at)}",
            "order": "usage_bucket_date.desc",
            "limit": "5000",
        },
    )

    relevant_events = [
        row
        for row in event_rows
        if (row.get("reconciliation_status") or "") != "reconciled_exact"
    ]
    relevant_buckets = [
        row
        for row in cost_bucket_rows
        if iso_day(starting_at) <= row.get("bucket_date", "")
    ]

    reconciled_by_event, reconciled_by_run = allocate_actual_costs(
        relevant_events,
        relevant_buckets,
    )

    for event in relevant_events:
        event_id = event["id"]
        if event_id not in reconciled_by_event:
            continue
        supabase.update_by_id(
            "provider_usage_events",
            event_id,
            {
                "reconciled_cost_usd": round(reconciled_by_event[event_id], 6),
                "reconciliation_status": "allocated_actual",
            },
        )

    run_rows = supabase.select(
        "generation_runs",
        "id",
        {
            "created_at": f"gte.{starting_at.isoformat()}",
            "order": "created_at.desc",
            "limit": "5000",
        },
    )

    for run in run_rows:
        run_id = run["id"]
        if run_id not in reconciled_by_run:
            continue
        supabase.update_by_id(
            "generation_runs",
            run_id,
            {
                "reconciled_cost_usd": round(reconciled_by_run[run_id], 6),
                "cost_reconciliation_status": "allocated_actual",
            },
        )

    print(
        json.dumps(
            {
                "days_back": days_back,
                "usage_buckets_upserted": len(usage_bucket_rows),
                "cost_buckets_upserted": len(cost_bucket_rows),
                "provider_events_reconciled": len(reconciled_by_event),
                "generation_runs_updated": len(reconciled_by_run),
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
