# 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.
"""Config file loading and llama-server presets.ini parsing."""

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

from . import providers as providers_mod

CONFIG_PATH = Path(
    os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config")
) / "llamachat" / "config.toml"

DEFAULTS = {
    "base_url": "http://localhost:8181",
    "presets": "/etc/llama-server/presets.ini",
    "socket": "",  # empty -> $XDG_RUNTIME_DIR/llamachat.sock
    "db": "",      # empty -> $XDG_DATA_HOME/llamachat/history.db
    "default_model": "",
    "request_timeout": 300,
    # Fraction of a model's context we allow a single attachment to fill.
    "attach_ctx_fraction": 0.5,
    # Rough chars-per-token used to turn a ctx-size into a char budget.
    "chars_per_token": 3.5,
    # Named prompt selected for new conversations. Empty means the global
    # default.md, and "none" means no system prompt at all.
    "default_prompt": "",
    # Web search through a SearXNG instance. Off unless both the flag is set
    # and a URL is given, so an upgrade never starts talking to the network.
    "search_enabled": False,
    "search_url": "",
    "search_results": 5,
    "search_snippet_chars": 300,
    "search_timeout": 10,
    # Searches per turn, after which the tool is withdrawn and the model
    # must answer. Two is safe now that the final round's tool result says
    # the tool is gone: without that note the model asks for a search it can
    # no longer make and the reply comes back empty. A quantized KV cache
    # (llama-server's cache-type-k/v) makes that failure far more likely, so
    # lower this to 1 if the server runs one.
    "max_searches": 2,
}

# The global prompt lives here; every other .md beside it is a named preset.
GLOBAL_PROMPT = "default"


@dataclass
class Preset:
    """One [section] of presets.ini, reduced to what the UI needs."""
    name: str
    vision: bool
    ctx_size: int

    def char_budget(self, fraction: float, chars_per_token: float) -> int:
        return int(self.ctx_size * chars_per_token * fraction)


@dataclass
class Config:
    base_url: str
    presets_path: Path
    socket_path: Path
    db_path: Path
    default_model: str
    request_timeout: int
    attach_ctx_fraction: float
    chars_per_token: float
    default_prompt: str
    prompts_dir: Path
    state_path: Path
    search_enabled: bool
    search_url: str
    search_results: int
    search_snippet_chars: int
    search_timeout: int
    max_searches: int
    providers: dict[str, providers_mod.Provider]
    models_path: Path
    provider_warnings: list[str]


def _runtime_dir() -> Path:
    rd = os.environ.get("XDG_RUNTIME_DIR")
    if rd:
        return Path(rd)
    # ponytail: /tmp fallback only matters on a broken session; no uid-dir
    # creation dance, XDG_RUNTIME_DIR is set on any real Wayland login.
    return Path("/tmp")


def _data_dir() -> Path:
    return Path(
        os.environ.get("XDG_DATA_HOME", Path.home() / ".local" / "share")
    ) / "llamachat"


def load(path: Path = CONFIG_PATH) -> Config:
    """Read config.toml, falling back to DEFAULTS for anything absent."""
    values = dict(DEFAULTS)
    if path.exists():
        with open(path, "rb") as fh:
            values.update(tomllib.load(fh))

    socket = values["socket"] or (_runtime_dir() / "llamachat.sock")
    db = values["db"] or (_data_dir() / "history.db")

    # A flag with no URL behind it is not enabled, it is misconfigured.
    # Resolving that here means nothing downstream has to test both.
    search_url = str(values["search_url"]).rstrip("/")
    search_enabled = bool(values["search_enabled"]) and bool(search_url)

    # Providers are built from the raw values so a bare base_url still
    # synthesizes the local entry. DEFAULTS supplies base_url when the file
    # names neither, which keeps a config with no network settings working.
    #
    # That default is load-bearing, not a convenience: it is the only reason
    # a malformed [providers.local] still leaves a working local provider,
    # since parse() rebuilds local from the bare URL after skipping it.
    # Dropping base_url from DEFAULTS, or letting a caller reach parse()
    # without one, would turn every such config into an app with no local
    # provider at all. Keep the key.
    provider_warnings: list[str] = []
    provider_table = providers_mod.parse(values, warnings=provider_warnings)

    return Config(
        base_url=str(values["base_url"]).rstrip("/"),
        presets_path=Path(values["presets"]).expanduser(),
        socket_path=Path(str(socket)).expanduser(),
        db_path=Path(str(db)).expanduser(),
        default_model=str(values["default_model"]),
        request_timeout=int(values["request_timeout"]),
        attach_ctx_fraction=float(values["attach_ctx_fraction"]),
        chars_per_token=float(values["chars_per_token"]),
        default_prompt=str(values["default_prompt"]),
        prompts_dir=path.parent / "prompts",
        # Window layout, remembered between runs. Not user-editable config,
        # so it sits beside it rather than inside config.toml.
        state_path=path.parent / "state.ini",
        search_enabled=search_enabled,
        search_url=search_url,
        search_results=int(values["search_results"]),
        search_snippet_chars=int(values["search_snippet_chars"]),
        search_timeout=int(values["search_timeout"]),
        max_searches=int(values["max_searches"]),
        providers=provider_table,
        # Cloud model metadata, cached beside state.ini for the same reason:
        # it is machine-written, not user-editable config.
        models_path=path.parent / "models.ini",
        provider_warnings=provider_warnings,
    )


def write_default(path: Path = CONFIG_PATH) -> Path:
    """Create a commented config.toml if none exists. Returns the path."""
    if path.exists():
        return path
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        "# llamachat configuration\n"
        f'base_url = "{DEFAULTS["base_url"]}"\n'
        f'presets = "{DEFAULTS["presets"]}"\n'
        '\n'
        '# Leave empty for XDG defaults:\n'
        '#   socket -> $XDG_RUNTIME_DIR/llamachat.sock\n'
        '#   db     -> $XDG_DATA_HOME/llamachat/history.db\n'
        'socket = ""\n'
        'db = ""\n'
        '\n'
        '# Model selected at startup. Empty picks the first available.\n'
        'default_model = ""\n'
        '\n'
        '# System prompt for new conversations. Prompts are markdown files\n'
        '# in the prompts/ directory beside this file; default.md is the\n'
        '# global one. Name a preset to use it instead, or "none" for no\n'
        '# system prompt. Empty means default.md.\n'
        'default_prompt = ""\n'
        '\n'
        '# Seconds before a generation request is abandoned.\n'
        f'request_timeout = {DEFAULTS["request_timeout"]}\n'
        '\n'
        '# How much of the model context one attachment may fill, and the\n'
        '# chars-per-token estimate used to convert ctx-size into characters.\n'
        f'attach_ctx_fraction = {DEFAULTS["attach_ctx_fraction"]}\n'
        f'chars_per_token = {DEFAULTS["chars_per_token"]}\n'
        '\n'
        '# Web search through SearXNG. The model decides when to search; it\n'
        '# is offered a web_search tool and calls it when a question needs\n'
        '# current information. Both search_enabled and search_url are\n'
        '# required: a flag with no URL stays off.\n'
        '#\n'
        '# The instance must have its JSON API enabled (format=json), which\n'
        '# many installs turn off by default.\n'
        '#\n'
        '# Search results are untrusted text entering the model context.\n'
        '# Only title, url and snippet are kept, snippets are truncated, and\n'
        '# result links are shown but never fetched automatically.\n'
        f'search_enabled = {str(DEFAULTS["search_enabled"]).lower()}\n'
        f'search_url = "{DEFAULTS["search_url"]}"\n'
        '\n'
        '# Results per query, characters kept from each snippet, seconds\n'
        '# before a query is abandoned, and searches allowed per turn.\n'
        '#\n'
        '# Lower max_searches to 1 if replies come back empty after a\n'
        '# search. That happens when the model asks for a search it can no\n'
        '# longer make, and a quantized KV cache on the server\n'
        '# (cache-type-k / cache-type-v) makes it far more likely.\n'
        f'search_results = {DEFAULTS["search_results"]}\n'
        f'search_snippet_chars = {DEFAULTS["search_snippet_chars"]}\n'
        f'search_timeout = {DEFAULTS["search_timeout"]}\n'
        f'max_searches = {DEFAULTS["max_searches"]}\n'
        '\n'
        '# External providers. The local router is a provider named "local",\n'
        '# synthesized from base_url above when no [providers.local] exists.\n'
        '# Any OpenAI-compatible endpoint works.\n'
        '#\n'
        '# Cloud models appear in the picker as provider:model. Local ones\n'
        '# stay bare, so nothing about the local setup changes.\n'
        '#\n'
        '# api_key accepts three forms:\n'
        '#   "pass:api/together"   read from the password store (preferred)\n'
        '#   "env:TOGETHER_KEY"    read from the environment\n'
        '#   "sk-..."              the key itself, in this file\n'
        '# It is read lazily, on the first request to that provider, so a\n'
        '# local-only session never unlocks the password store.\n'
        '#\n'
        '# filter keeps only models whose id contains one of these strings,\n'
        '# case-insensitively. Providers list hundreds of models; without a\n'
        '# filter the picker is unusable. Omit it to list them all.\n'
        '#\n'
        '# ctx_size, vision, price_in and price_out prefill the per-model\n'
        '# dialog. Prices are US dollars per million tokens. Everything the\n'
        '# dialog saves goes to models.ini beside this file, so none of\n'
        '# these has to be set here.\n'
        '#\n'
        '# [providers.together]\n'
        '# base_url = "https://api.together.xyz"\n'
        '# api_key = "pass:api/together"\n'
        '# filter = ["qwen", "deepseek"]\n'
        '# ctx_size = 32768\n'
        '# price_in = 0.60\n'
        '# price_out = 0.60\n'
    )
    return path


def parse_presets(path: Path) -> dict[str, Preset]:
    """Parse llama-server's presets.ini.

    Both ';' and '#' start a comment there, so a commented-out `#mmproj`
    line correctly reads as "this preset has no vision support".
    """
    parser = configparser.ConfigParser(
        comment_prefixes=(";", "#"),
        inline_comment_prefixes=(";", "#"),
        strict=False,
    )
    try:
        text = path.read_text(encoding="utf-8")
    except OSError:
        return {}

    # presets.ini opens with a bare `version = 1` before any section, which
    # configparser rejects. Give those leading keys a section to live in.
    try:
        parser.read_string("[__file__]\n" + text)
    except configparser.Error:
        return {}

    presets: dict[str, Preset] = {}
    for name in parser.sections():
        if name == "__file__":
            continue
        section = parser[name]
        try:
            ctx = int(section.get("ctx-size", "4096"))
        except ValueError:
            ctx = 4096
        presets[name] = Preset(
            name=name,
            vision=bool(section.get("mmproj", "").strip()),
            ctx_size=ctx,
        )
    return presets
