from __future__ import annotations

import threading
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from decimal import Decimal
from typing import Any, Callable, Optional

dict_row: Callable[..., Any] | None = None
ConnectionPool: type | None = None
try:  # pragma: no cover
    from psycopg.rows import dict_row as _dict_row
except ImportError:  # pragma: no cover
    pass
else:
    dict_row = _dict_row

try:  # pragma: no cover
    from psycopg_pool import ConnectionPool as _ConnectionPool
except ImportError:  # pragma: no cover
    pass
else:
    ConnectionPool = _ConnectionPool


@dataclass
class FBR:
    """Finalized Billing Record (FBR).

    An FBR represents a complete billing transaction for a single API request.
    It tracks token usage, costs, and billing metadata. FBRs are keyed by
    (org_id, idempotency_key) to support idempotent billing.

    In P0, FBRs are stored in-memory or Postgres. Production systems should
    use durable storage with proper reconciliation.

    Attributes:
        id: Unique identifier for this billing record
        org_id: Organization being billed
        idempotency_key: Deduplication key
        status: Record status (pending|finalized)
        price_book_version: Version of pricing data used
        billing_mode: How to bill (resell|passthrough)
        prompt_tokens: Number of input tokens consumed
        completion_tokens: Number of output tokens generated
        cost_usd: Total cost in USD
        api_key_id: API key that authorized this request
        finalized_at: When the record was finalized (ISO timestamp)
        created_at: When the record was created (ISO timestamp)
    """

    id: str
    org_id: str
    idempotency_key: str
    status: str = "pending"
    price_book_version: str = "p0-0"
    billing_mode: str = "resell"
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cost_usd: float = 0.0
    api_key_id: Optional[str] = None
    finalized_at: Optional[str] = None
    created_at: Optional[str] = None


class IdempotencyStore(ABC):
    @abstractmethod
    def claim(self, org_id: str, idempotency_key: str, request_id: str) -> bool:  # noqa: D401
        """Attempt to claim an idempotency key for a request.

        Returns True if claimed, False if another request already owns it.
        """

    @abstractmethod
    def get_request_id(self, org_id: str, idempotency_key: str) -> Optional[str]:
        """Return existing request id for the idempotency key if present."""


class InMemoryIdempotencyStore(IdempotencyStore):
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._records: dict[tuple[str, str], str] = {}  # (org_id, idem_key) -> request_id

    def claim(self, org_id: str, idempotency_key: str, request_id: str) -> bool:
        with self._lock:
            key = (org_id, idempotency_key)
            if key in self._records and self._records[key] != request_id:
                return False
            self._records[key] = request_id
            return True

    def get_request_id(self, org_id: str, idempotency_key: str) -> Optional[str]:
        with self._lock:
            return self._records.get((org_id, idempotency_key))


class BillingStore(ABC):
    @abstractmethod
    def seed_fbr(
        self,
        fbr_id: str,
        org_id: str,
        idempotency_key: str,
        price_book_version: str,
        *,
        billing_mode: str = "resell",
        api_key_id: Optional[str] = None,
    ) -> FBR:
        pass

    @abstractmethod
    def record_usage_delta(
        self, fbr_id: str, prompt_tokens: int, completion_tokens: int, cost_usd: float
    ) -> None:
        pass

    @abstractmethod
    def finalize(self, fbr_id: str) -> Optional[FBR]:
        pass

    @abstractmethod
    def reconcile_pending(self, *, older_than_minutes: int = 30) -> int:
        """Finalize pending billing records older than the threshold."""

    # Optional streaming-friendly API; default no-op if not implemented
    def record_usage_event(
        self,
        fbr_id: str,
        event_seq: int,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> None:
        # Default: accumulate via delta if an FBR exists bound to request_id.
        # The in-memory implementation overrides this to aggregate per-event.
        return None

    def list_records(
        self, org_id: str, *, limit: int = 50, cursor: Optional[str] = None
    ) -> list[dict[str, Any]]:
        return []

    def get_org_plan(self, org_id: str) -> Optional[dict[str, Any]]:
        return None


class InMemoryBillingStore(BillingStore):
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._fbr_by_key: dict[tuple[str, str], FBR] = {}
        self._fbr_by_id: dict[str, FBR] = {}
        self._events: dict[str, set[int]] = {}  # request_id -> seen event ids

    def seed_fbr(
        self,
        fbr_id: str,
        org_id: str,
        idempotency_key: str,
        price_book_version: str,
        *,
        billing_mode: str = "resell",
        api_key_id: Optional[str] = None,
    ) -> FBR:
        with self._lock:
            key = (org_id, idempotency_key)
            if key in self._fbr_by_key:
                return self._fbr_by_key[key]
            fbr = FBR(
                id=fbr_id,
                org_id=org_id,
                idempotency_key=idempotency_key,
                price_book_version=price_book_version,
                billing_mode=billing_mode,
                api_key_id=api_key_id,
                created_at=datetime.now(UTC).isoformat(),
            )
            self._fbr_by_key[key] = fbr
            self._fbr_by_id[fbr_id] = fbr
            return fbr

    def record_usage_delta(
        self, fbr_id: str, prompt_tokens: int, completion_tokens: int, cost_usd: float
    ) -> None:
        with self._lock:
            fbr = self._fbr_by_id.get(fbr_id)
            if not fbr:
                return
            fbr.prompt_tokens += int(prompt_tokens)
            fbr.completion_tokens += int(completion_tokens)
            fbr.cost_usd += float(cost_usd)

    def record_usage_event(
        self,
        fbr_id: str,
        event_seq: int,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> None:
        with self._lock:
            seen = self._events.setdefault(fbr_id, set())
            if event_seq in seen:
                return
            seen.add(event_seq)
            fbr = self._fbr_by_id.get(fbr_id)
            if not fbr:
                return
            fbr.prompt_tokens += int(prompt_tokens)
            fbr.completion_tokens += int(completion_tokens)
            fbr.cost_usd += float(cost_usd)

    def finalize(self, fbr_id: str) -> Optional[FBR]:
        with self._lock:
            fbr = self._fbr_by_id.get(fbr_id)
            if not fbr:
                return None
            fbr.status = "finalized"
            fbr.finalized_at = datetime.now(UTC).isoformat()
            return fbr

    def reconcile_pending(self, *, older_than_minutes: int = 30) -> int:
        cutoff = datetime.now(UTC) - timedelta(minutes=older_than_minutes)
        finalized_at = datetime.now(UTC).isoformat()
        updated = 0
        with self._lock:
            for fbr in self._fbr_by_id.values():
                if fbr.status != "pending":
                    continue
                if fbr.created_at is None:
                    raise ValueError("In-memory billing record missing created_at")
                created_at = datetime.fromisoformat(fbr.created_at)
                if created_at.tzinfo is None:
                    created_at = created_at.replace(tzinfo=UTC)
                if created_at < cutoff:
                    fbr.status = "finalized"
                    fbr.finalized_at = finalized_at
                    updated += 1
        return updated


# Placeholders for Postgres/Redis-backed stores (production drivers)


class PostgresBillingStore(BillingStore):
    """Production billing store backed by Postgres."""

    def __init__(
        self,
        dsn: str,
        *,
        table: str = "billing_records",
        usage_table: str = "billing_usage_events",
        auto_init: bool = False,
        pool_min_size: int = 1,
        pool_max_size: int = 10,
    ) -> None:
        if ConnectionPool is None:
            raise ImportError("psycopg must be installed to use PostgresBillingStore")
        self.table = table
        self.usage_table = usage_table
        self._pool = ConnectionPool(conninfo=dsn, min_size=pool_min_size, max_size=pool_max_size)
        if auto_init:
            self._auto_init()

    def close(self) -> None:
        if self._pool:
            self._pool.close()

    def _auto_init(self) -> None:
        ddl = [
            f"""
            CREATE TABLE IF NOT EXISTS {self.table} (
                id TEXT PRIMARY KEY,
                org_id TEXT NOT NULL,
                api_key_id UUID,
                idempotency_key TEXT NOT NULL,
                status billing_status NOT NULL DEFAULT 'pending',
                price_book_version TEXT NOT NULL,
                billing_mode billing_mode NOT NULL DEFAULT 'resell',
                prompt_tokens_total BIGINT NOT NULL DEFAULT 0,
                completion_tokens_total BIGINT NOT NULL DEFAULT 0,
                cost_usd_total NUMERIC(18,6) NOT NULL DEFAULT 0,
                created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
                finalized_at TIMESTAMPTZ,
                metadata JSONB NOT NULL DEFAULT '{{}}'::jsonb
            );
            """,
            (
                f"CREATE UNIQUE INDEX IF NOT EXISTS {self.table}_org_idem_idx "
                f"ON {self.table}(org_id, idempotency_key);"
            ),
            f"""
            CREATE TABLE IF NOT EXISTS {self.usage_table} (
                id BIGSERIAL PRIMARY KEY,
                fbr_id TEXT NOT NULL REFERENCES {self.table}(id) ON DELETE CASCADE,
                event_seq INT NOT NULL,
                prompt_tokens_delta INT NOT NULL,
                completion_tokens_delta INT NOT NULL,
                cost_usd_delta NUMERIC(18,6) NOT NULL,
                created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
                UNIQUE (fbr_id, event_seq)
            );
            """,
            f"CREATE INDEX IF NOT EXISTS {self.usage_table}_fbr_idx ON {self.usage_table}(fbr_id);",
        ]
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                for stmt in ddl:
                    cur.execute(stmt)
            conn.commit()

    def seed_fbr(
        self,
        fbr_id: str,
        org_id: str,
        idempotency_key: str,
        price_book_version: str,
        *,
        billing_mode: str = "resell",
        api_key_id: Optional[str] = None,
    ) -> FBR:
        with self._pool.connection() as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    f"""
                    INSERT INTO {self.table} (
                        id,
                        org_id,
                        api_key_id,
                        idempotency_key,
                        price_book_version,
                        billing_mode
                    )
                    VALUES (%s,%s,%s,%s,%s,%s)
                    ON CONFLICT (org_id, idempotency_key) DO NOTHING
                    RETURNING id, org_id, api_key_id, idempotency_key, status, price_book_version,
                              billing_mode, prompt_tokens_total, completion_tokens_total,
                              cost_usd_total, finalized_at, created_at
                    """,
                    (fbr_id, org_id, api_key_id, idempotency_key, price_book_version, billing_mode),
                )
                row = cur.fetchone()
                if not row:
                    cur.execute(
                        f"""
                        SELECT id, org_id, api_key_id, idempotency_key, status, price_book_version,
                               billing_mode, prompt_tokens_total, completion_tokens_total,
                               cost_usd_total, finalized_at, created_at
                        FROM {self.table}
                        WHERE org_id = %s AND idempotency_key = %s
                        """,
                        (org_id, idempotency_key),
                    )
                    row = cur.fetchone()
            conn.commit()
        return _row_to_fbr(row)

    def record_usage_event(
        self,
        fbr_id: str,
        event_seq: int,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> None:
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"""
                    INSERT INTO {self.usage_table} (
                        fbr_id,
                        event_seq,
                        prompt_tokens_delta,
                        completion_tokens_delta,
                        cost_usd_delta
                    )
                    VALUES (%s,%s,%s,%s,%s)
                    ON CONFLICT (fbr_id, event_seq) DO NOTHING
                    """,
                    (
                        fbr_id,
                        int(event_seq),
                        int(prompt_tokens),
                        int(completion_tokens),
                        Decimal(str(cost_usd)),
                    ),
                )
                if cur.rowcount:
                    cur.execute(
                        f"""
                        UPDATE {self.table}
                        SET prompt_tokens_total = prompt_tokens_total + %s,
                            completion_tokens_total = completion_tokens_total + %s,
                            cost_usd_total = cost_usd_total + %s
                        WHERE id = %s AND status = 'pending'
                        """,
                        (
                            int(prompt_tokens),
                            int(completion_tokens),
                            Decimal(str(cost_usd)),
                            fbr_id,
                        ),
                    )
            conn.commit()

    def record_usage_delta(
        self, fbr_id: str, prompt_tokens: int, completion_tokens: int, cost_usd: float
    ) -> None:
        self.record_usage_event(fbr_id, 0, prompt_tokens, completion_tokens, cost_usd)

    def finalize(self, fbr_id: str) -> Optional[FBR]:
        with self._pool.connection() as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    f"""
                    UPDATE {self.table}
                    SET status='finalized', finalized_at = COALESCE(finalized_at, now())
                    WHERE id = %s
                    RETURNING id, org_id, api_key_id, idempotency_key, status, price_book_version,
                              billing_mode, prompt_tokens_total, completion_tokens_total,
                              cost_usd_total, finalized_at, created_at
                    """,
                    (fbr_id,),
                )
                row = cur.fetchone()
                if not row:
                    cur.execute(
                        f"""
                        SELECT id, org_id, api_key_id, idempotency_key, status, price_book_version,
                               billing_mode, prompt_tokens_total, completion_tokens_total,
                               cost_usd_total, finalized_at, created_at
                        FROM {self.table}
                        WHERE id = %s
                        """,
                        (fbr_id,),
                    )
                    row = cur.fetchone()
            conn.commit()
        return _row_to_fbr(row) if row else None

    def get_org_plan(self, org_id: str) -> Optional[dict[str, Any]]:
        with self._pool.connection() as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        plan_id,
                        tier,
                        price_book_version,
                        default_billing_mode,
                        provider_overrides
                    FROM org_plans
                    WHERE org_id = %s
                    """,
                    (org_id,),
                )
                row = cur.fetchone()
        if not row:
            return None
        data = dict(row)
        overrides = data.get("provider_overrides")
        if not isinstance(overrides, dict):
            data["provider_overrides"] = {}
        return data

    def list_records(
        self, org_id: str, *, limit: int = 50, cursor: Optional[str] = None
    ) -> list[dict[str, Any]]:
        sql = (
            "SELECT id, status, price_book_version, billing_mode, prompt_tokens_total, "
            "completion_tokens_total, cost_usd_total, created_at, finalized_at "
            f"FROM {self.table} WHERE org_id = %s"
        )
        params: list[Any] = [org_id]
        if cursor:
            sql += " AND created_at < %s"
            params.append(_parse_timestamp(cursor))
        sql += " ORDER BY created_at DESC LIMIT %s"
        params.append(limit)

        with self._pool.connection() as conn:
            with conn.cursor(row_factory=dict_row) as cur:
                cur.execute(sql, tuple(params))
                rows = cur.fetchall()
        for row in rows:
            for field in ("created_at", "finalized_at"):
                if row.get(field) is not None:
                    row[field] = row[field].isoformat()
            for field in ("prompt_tokens_total", "completion_tokens_total"):
                value = row.get(field)
                if value is not None:
                    row[field] = int(value)
            cost = row.get("cost_usd_total")
            if isinstance(cost, Decimal):
                row["cost_usd_total"] = float(cost)
        return [dict(row) for row in rows]

    def reconcile_pending(self, *, older_than_minutes: int = 30) -> int:
        cutoff = datetime.now(UTC) - timedelta(minutes=older_than_minutes)
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"""
                    UPDATE {self.table}
                    SET status='finalized', finalized_at = COALESCE(finalized_at, now())
                    WHERE status='pending' AND created_at < %s
                    RETURNING id
                    """,
                    (cutoff,),
                )
                rows = cur.fetchall()
            conn.commit()
        return len(rows)


class RedisIdempotencyStore(IdempotencyStore):
    """Production idempotency store (Redis) placeholder using SET NX + TTL."""

    def __init__(self, client: Any, ttl_seconds: int = 24 * 3600) -> None:
        self.client = client
        self.ttl = ttl_seconds

    def claim(self, org_id: str, idempotency_key: str, request_id: str) -> bool:
        key = f"idem:{org_id}:{idempotency_key}"
        ok = self.client.set(key, request_id, nx=True, ex=self.ttl)
        return bool(ok)

    def get_request_id(self, org_id: str, idempotency_key: str) -> Optional[str]:
        key = f"idem:{org_id}:{idempotency_key}"
        val = self.client.get(key)
        if val is None:
            return None
        if isinstance(val, bytes):
            return val.decode("utf-8")
        return str(val)


def _row_to_fbr(row: Optional[dict[str, Any]]) -> FBR:
    if row is None:
        raise ValueError("Billing record row is required")
    record_id = row.get("id")
    if record_id is None:
        raise ValueError("Billing record missing required 'id' field")
    org_id = row.get("org_id")
    if org_id is None:
        raise ValueError("Billing record missing required 'org_id' field")
    cost = row.get("cost_usd_total", 0)
    if isinstance(cost, Decimal):
        cost = float(cost)
    finalized_at = row.get("finalized_at")
    if finalized_at is not None and hasattr(finalized_at, "isoformat"):
        finalized_at = finalized_at.isoformat()
    created_at = row.get("created_at")
    if created_at is not None and hasattr(created_at, "isoformat"):
        created_at = created_at.isoformat()
    api_key_id = row.get("api_key_id")
    if api_key_id is not None:
        api_key_id = str(api_key_id)
    return FBR(
        id=str(record_id),
        org_id=str(org_id),
        idempotency_key=row.get("idempotency_key", ""),
        status=row.get("status", "pending"),
        price_book_version=row.get("price_book_version", "p0-0"),
        billing_mode=row.get("billing_mode", "resell"),
        prompt_tokens=int(row.get("prompt_tokens_total", 0) or 0),
        completion_tokens=int(row.get("completion_tokens_total", 0) or 0),
        cost_usd=float(cost or 0.0),
        api_key_id=api_key_id,
        finalized_at=finalized_at,
        created_at=created_at,
    )


def _parse_timestamp(value: str) -> datetime:
    return datetime.fromisoformat(value)
