# SPDX-License-Identifier: GPL-2.0-only
#
# llamachat - a small native chat client for a local llama.cpp router
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
"""Per-model metadata: what presets.ini cannot answer for a cloud model."""

import configparser
import math
import os
from dataclasses import dataclass
from pathlib import Path

from . import providers as providers_mod

# configparser reads a literal [DEFAULT] as inherited defaults for every
# other section. A model genuinely named DEFAULT is stored under this name
# instead, so the file stays safe for the stock-configparser readers that
# Task 15 invites by documenting it for hand-editing.
#
# ponytail: only the exact id "DEFAULT" collides. A file already containing
# a hand-written [DEFAULT] is handled by _empty() below, not here, and no
# other [DEFAULT]-alike is defended against. Provider ids arrive as
# "provider:model", so a bare DEFAULT can only be a local model.
_DEFAULT_ID = "DEFAULT"
_DEFAULT_ESCAPED = "\x00llamachat-model-DEFAULT"

# Recorded when the dialog is cancelled, so a model tried once never nags.
# Only the section's presence is read, never this value; the key exists to
# give an otherwise empty section something readable in the file. It means
# "no real keys here", and mark_skipped() is its only writer.
SKIPPED = "offered"


@dataclass
class ModelInfo:
    """What a model can tell us. Every field may be unknown."""

    ctx_size: int | None = None
    vision: bool | None = None
    price_in: float | None = None
    price_out: float | None = None

    def is_empty(self) -> bool:
        return all(
            v is None
            for v in (self.ctx_size, self.vision, self.price_in, self.price_out)
        )


def _get(section, key, cast):
    raw = section.get(key, "").strip()
    if not raw:
        return None
    try:
        value = cast(raw)
    except ValueError:
        # A hand-edited "32k" or "3.7" for an int degrades to unknown rather
        # than raising: the file is user-editable, and one bad line must not
        # take down every other model's metadata with it.
        return None
    # float() accepts "nan" and "inf", which would survive every later check
    # and surface as a "$nan" cost. Caught here rather than at each cost call
    # site, so the arithmetic downstream can assume a real number.
    if cast is float and not math.isfinite(value):
        return None
    return value


def _get_bool(section, key):
    raw = section.get(key, "").strip().lower()
    if raw in ("true", "yes", "1", "on"):
        return True
    if raw in ("false", "no", "0", "off"):
        return False
    return None


class ModelStore:
    """models.ini, keyed by full model id.

    configparser rather than TOML because this file is written by the app,
    and tomllib is read-only in the standard library.
    """

    def __init__(self, path: Path):
        self.path = Path(path)
        # Model ids contain colons and slashes, so no key/value delimiter
        # may be inferred from a section name. Sections are safe as-is:
        # brackets, colons, equals and percent signs all round-trip, and
        # interpolation is off so a "%" in an id or value is literal.
        self.parser = self._empty()
        if self.path.exists():
            try:
                self.parser.read(self.path, encoding="utf-8")
            except (configparser.Error, OSError, UnicodeDecodeError):
                # ponytail: a corrupt file reads as empty; the dialog can
                # rewrite it. Failing to start over metadata is worse. The
                # cost is that the next save silently discards whatever was
                # readable, which is acceptable for a cache of prices.
                self.parser = self._empty()

    @staticmethod
    def _empty():
        # default_section is renamed away from "DEFAULT" so that a hand-added
        # [DEFAULT] in the file is read as an ordinary section rather than as
        # defaults inherited by every other one. Without this, one such block
        # silently gives every model its ctx_size and prices, including
        # models with no entry at all. Reserving a name containing a NUL,
        # which no model id can hold, keeps every section independent.
        return configparser.ConfigParser(
            interpolation=None, default_section="\x00llamachat-default"
        )

    @staticmethod
    def _key(model_id: str) -> str:
        """The section name for a model id. Only 'DEFAULT' is rewritten."""
        return _DEFAULT_ESCAPED if model_id == _DEFAULT_ID else model_id

    def get(self, model_id: str) -> ModelInfo | None:
        """Stored metadata, or None when there is none worth having.

        None and was_offered() True together mean "asked, learned nothing":
        either the dialog was cancelled or every field was left blank. The
        caller wants exactly that distinction, so it is not a trap: get()
        answers "what do we know", was_offered() answers "should we ask".
        """
        name = self._key(model_id)
        if not self.parser.has_section(name):
            return None
        section = self.parser[name]
        info = ModelInfo(
            ctx_size=_get(section, "ctx_size", int),
            vision=_get_bool(section, "vision"),
            price_in=_get(section, "price_in", float),
            price_out=_get(section, "price_out", float),
        )
        return None if info.is_empty() else info

    def was_offered(self, model_id: str) -> bool:
        """Whether the dialog has already been shown for this model."""
        return self.parser.has_section(self._key(model_id))

    def save(self, model_id: str, info: ModelInfo) -> None:
        section = self._section(model_id)
        for key, value in (
            ("ctx_size", info.ctx_size),
            ("vision", info.vision),
            ("price_in", info.price_in),
            ("price_out", info.price_out),
        ):
            if value is None:
                section.pop(key, None)
            elif isinstance(value, bool):
                # Before the numeric branch: bool is a subclass of int, so
                # str(False) would otherwise write "False" for vision.
                section[key] = "true" if value else "false"
            else:
                section[key] = str(value)
        # The marker means "no real keys", and mark_skipped() is its only
        # writer. Dropping it here keeps that invariant with one enforcer
        # instead of two that have to agree. A save with every field blank
        # leaves an empty section, which configparser round-trips fine.
        section.pop(SKIPPED, None)
        self._write()

    def mark_skipped(self, model_id: str) -> None:
        """Record a cancelled dialog: offered, declined, do not ask again.

        Never clears fields already stored. Cancelling a dialog opened over
        a model we know about must not delete what we know.
        """
        section = self._section(model_id)
        if not any(k != SKIPPED for k in section):
            section[SKIPPED] = "true"
        self._write()

    def _section(self, model_id: str):
        name = self._key(model_id)
        if not self.parser.has_section(name):
            self.parser.add_section(name)
        return self.parser[name]

    def _write(self) -> None:
        """Rewrite the whole file from the in-memory parser.

        ponytail: last writer wins. Each store holds a full parser, so two
        live instances would drop each other's sections rather than merge.
        Not defended against because ipc.py enforces a single instance and
        the window holds exactly one store, making a second writer a case
        that cannot arise. The upgrade path, if the app ever grows a second
        writer, is to reread before each save and merge.
        """
        self.path.parent.mkdir(parents=True, exist_ok=True)
        # Written to a sibling and renamed: os.replace is atomic within a
        # directory, so a crash mid-write leaves the old file intact instead
        # of a truncated one. Cheap here, and the alternative is losing every
        # recorded price to one bad moment.
        tmp = self.path.with_name(self.path.name + f".{os.getpid()}.tmp")
        try:
            with open(tmp, "w", encoding="utf-8") as fh:
                self.parser.write(fh)
            os.replace(tmp, self.path)
        except OSError:
            tmp.unlink(missing_ok=True)
            raise


def resolve(model_id: str, table: dict, store: "ModelStore") -> ModelInfo:
    """Merge the three metadata layers, most specific first.

    Per field, not per source: a models.ini entry that sets only ctx_size
    still inherits the provider's prices.

    These three are not the whole merge. Task 12's model_info() in ui.py
    layers presets.ini on top of what this returns, because models.py has
    no business importing presets. The consequence of that split is that
    "what does this model's ctx_size resolve to" has no single answer site,
    so a reader here is looking at three of four layers.
    """
    stored = store.get(model_id) or ModelInfo()
    name, _ = providers_mod.split(model_id, table)
    provider = table.get(name)

    def pick(from_store, from_provider):
        return from_store if from_store is not None else from_provider

    if provider is None:
        return stored
    return ModelInfo(
        ctx_size=pick(stored.ctx_size, provider.ctx_size),
        vision=pick(stored.vision, provider.vision),
        price_in=pick(stored.price_in, provider.price_in),
        price_out=pick(stored.price_out, provider.price_out),
    )


def is_billable(model_id: str, table: dict) -> bool:
    """Whether this model costs money, regardless of prices being known.

    Keyed on the provider having an api_key: that is what distinguishes a
    free local model from a cloud one whose price was never entered.

    Alone among these functions it takes no store, and that asymmetry is
    load-bearing rather than an omission: whether a model costs money does
    not depend on which prices we happen to have recorded. Pairing it with
    is_priced is the point, one answers "is this billed", the other "can we
    say how much". Do not add a store parameter to make it match.
    """
    name, _ = providers_mod.split(model_id, table)
    provider = table.get(name)
    return bool(provider and not provider.is_local and provider.api_key)


def is_priced(model_id: str, table: dict, store: "ModelStore") -> bool:
    """Whether a cost can be computed for this model.

    A non-finite price is one no cost can be computed from, so it reads as
    unpriced. Both of today's writers filter these out before they get here,
    which makes this the predicate stated completely rather than a redundant
    check: answering True for a nan would send the readout down the priced
    branch to render "$nan" in the one state built to admit it cannot say.
    """
    info = resolve(model_id, table, store)
    return any(
        p is not None and math.isfinite(p) for p in (info.price_in, info.price_out)
    )


def message_cost(row, table: dict, store: "ModelStore") -> float:
    """Cost of one stored assistant row, priced at the model that made it.

    ponytail: binary floats for money, which is normally a mistake. It is
    not one here because these are estimates against hand-typed prices,
    never reconciled against a real bill, and Task 11's tooltip says
    "Approximate". The magnitudes involved sit nowhere near double's
    precision limit. The ceiling is that summing many rows accumulates
    representation error, so a long conversation's total can differ in the
    last cents from the same rows added another way. The upgrade path is
    integer micro-dollars end to end, not Decimal: token counts are already
    integers, so prices scaled to millionths keep the arithmetic exact and
    the storage an INTEGER column.
    """
    model_id = row["model"] if row["model"] else ""
    if not model_id:
        return 0.0
    info = resolve(model_id, table, store)
    prompt = row["prompt_tokens"] or 0
    completion = row["completion_tokens"] or 0
    total = 0.0
    if info.price_in is not None:
        total += prompt * info.price_in
    if info.price_out is not None:
        total += completion * info.price_out
    return total / 1_000_000


def conversation_cost(rows, table: dict, store: "ModelStore") -> float:
    """Everything spent in one conversation so far.

    Rows predating the token columns carry NULLs and contribute zero, so an
    old conversation reads as free rather than as a fabricated number.
    """
    return sum(message_cost(row, table, store) for row in rows)


def projected_cost(
    tokens: int, model_id: str, table: dict, store: "ModelStore"
) -> float:
    """What sending `tokens` of input to this model would cost.

    Only the input rate: the length of the reply is unknowable in advance.
    """
    info = resolve(model_id, table, store)
    if info.price_in is None:
        return 0.0
    return tokens * info.price_in / 1_000_000


def format_cost(amount: float) -> str:
    """Money, with enough precision to see a cheap turn move the number."""
    if amount < 1:
        return f"${amount:.3f}"
    return f"${amount:.2f}"
