# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
from datetime import datetime, timedelta, timezone
from typing import Any

from devplacepy.database import db, get_table

logger = logging.getLogger("devii.store")

CONVERSATIONS = "devii_conversations"
LEDGER = "devii_usage_ledger"
TURNS = "devii_turns"

PROMPT_AUDIT_CHARS = 4000
REPLY_AUDIT_CHARS = 4000


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _iso(moment: datetime) -> str:
    return moment.isoformat(timespec="microseconds")


class ConversationStore:
    def load(
        self, owner_kind: str, owner_id: str, channel: str = "main"
    ) -> list[dict[str, Any]] | None:
        if CONVERSATIONS not in db.tables:
            return None
        row = get_table(CONVERSATIONS).find_one(
            owner_kind=owner_kind, owner_id=owner_id, channel=channel
        )
        if not row or not row.get("messages"):
            return None
        try:
            return json.loads(row["messages"])
        except (ValueError, TypeError):
            logger.warning("Corrupt conversation for %s/%s", owner_kind, owner_id)
            return None

    def save(
        self,
        owner_kind: str,
        owner_id: str,
        messages: list[dict[str, Any]],
        channel: str = "main",
    ) -> None:
        now = _iso(_now())
        record = {
            "owner_kind": owner_kind,
            "owner_id": owner_id,
            "channel": channel,
            "messages": json.dumps(messages),
            "updated_at": now,
        }
        table = get_table(CONVERSATIONS)
        existing = (
            table.find_one(owner_kind=owner_kind, owner_id=owner_id, channel=channel)
            if CONVERSATIONS in db.tables
            else None
        )
        if existing:
            table.update({**record, "id": existing["id"]}, ["id"])
        else:
            record["created_at"] = now
            table.insert(record)

    def clear(self, owner_kind: str, owner_id: str, channel: str = "main") -> None:
        if CONVERSATIONS in db.tables:
            get_table(CONVERSATIONS).delete(
                owner_kind=owner_kind, owner_id=owner_id, channel=channel
            )


class UsageLedger:
    def record(
        self,
        owner_kind: str,
        owner_id: str,
        turn_id: str,
        usage: dict[str, int],
        cost_usd: float,
        model: str,
    ) -> None:
        get_table(LEDGER).insert(
            {
                "owner_kind": owner_kind,
                "owner_id": owner_id,
                "turn_id": turn_id,
                "created_at": _iso(_now()),
                "prompt_tokens": int(usage.get("prompt_tokens", 0)),
                "completion_tokens": int(usage.get("completion_tokens", 0)),
                "cache_hit_tokens": int(usage.get("cache_hit_tokens", 0)),
                "cache_miss_tokens": int(usage.get("cache_miss_tokens", 0)),
                "cost_usd": float(cost_usd),
                "model": model,
            }
        )

    def spent_24h(self, owner_kind: str, owner_id: str) -> float:
        if LEDGER not in db.tables:
            return 0.0
        cutoff = _iso(_now() - timedelta(hours=24))
        rows = db.query(
            f"SELECT COALESCE(SUM(cost_usd), 0) AS total FROM {LEDGER} "
            "WHERE owner_kind = :owner_kind AND owner_id = :owner_id AND created_at >= :cutoff",
            owner_kind=owner_kind,
            owner_id=owner_id,
            cutoff=cutoff,
        )
        for row in rows:
            return float(row["total"] or 0.0)
        return 0.0

    def turns_24h(self, owner_kind: str, owner_id: str) -> int:
        if LEDGER not in db.tables:
            return 0
        cutoff = _iso(_now() - timedelta(hours=24))
        rows = db.query(
            f"SELECT COUNT(*) AS n FROM {LEDGER} "
            "WHERE owner_kind = :owner_kind AND owner_id = :owner_id AND created_at >= :cutoff",
            owner_kind=owner_kind,
            owner_id=owner_id,
            cutoff=cutoff,
        )
        for row in rows:
            return int(row["n"] or 0)
        return 0

    def total_spent_24h(self) -> float:
        if LEDGER not in db.tables:
            return 0.0
        cutoff = _iso(_now() - timedelta(hours=24))
        rows = db.query(
            f"SELECT COALESCE(SUM(cost_usd), 0) AS total FROM {LEDGER} WHERE created_at >= :cutoff",
            cutoff=cutoff,
        )
        for row in rows:
            return float(row["total"] or 0.0)
        return 0.0

    def prune(self, older_than_hours: int = 48) -> int:
        if LEDGER not in db.tables:
            return 0
        cutoff = _iso(_now() - timedelta(hours=older_than_hours))
        return int(get_table(LEDGER).delete(created_at={"<": cutoff}))

    def reset(self, owner_kind: str, owner_id: str) -> int:
        if LEDGER not in db.tables:
            return 0
        table = get_table(LEDGER)
        count = table.count(owner_kind=owner_kind, owner_id=owner_id)
        table.delete(owner_kind=owner_kind, owner_id=owner_id)
        return int(count)

    def reset_owner_kind(self, owner_kind: str) -> int:
        if LEDGER not in db.tables:
            return 0
        table = get_table(LEDGER)
        count = table.count(owner_kind=owner_kind)
        table.delete(owner_kind=owner_kind)
        return int(count)

    def reset_all(self) -> int:
        if LEDGER not in db.tables:
            return 0
        table = get_table(LEDGER)
        count = table.count()
        table.delete()
        return int(count)


class TurnAudit:
    def record(self, record: dict[str, Any]) -> None:
        record = dict(record)
        record["prompt"] = (record.get("prompt") or "")[:PROMPT_AUDIT_CHARS]
        record["reply"] = (record.get("reply") or "")[:REPLY_AUDIT_CHARS]
        get_table(TURNS).insert(record)
