"""
Copyright (C) 2024-2026 Mo Zhou <lumin@debian.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

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 Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
import os
import re
import sys
from typing import Dict, Iterable, Optional

import urwid

from . import arguments, defaults

default = defaults.Config()
console = defaults.console

_TITLE = "DebGPT Configurator"


def _abort_on_None(value: Optional[str]) -> None:
    if value is None:
        print("Aborted.")
        exit(1)


class ListBoxVimKeys(urwid.ListBox):
    def keypress(self, size, key) -> None:
        key_map = {
            "k": "up",
            "j": "down",
        }
        if key in ("esc", "q", "Q"):
            sys.exit(1)
        super().keypress(size, key_map.get(key, key))


class SingleChoice(object):
    # we store the user choice here
    _choice: Optional[str] = None

    @staticmethod
    def exit_on_esc(key: str) -> None:
        if key in ("esc",):
            sys.exit(1)

    def item_chosen(self, button: urwid.Button, choice: str) -> None:
        SingleChoice._choice = choice
        raise urwid.ExitMainLoop(choice)

    def __init__(
        self,
        title: str,
        question: str,
        choices: Iterable[str],
        helpmsg: str,
        statusmsg: str,
        *,
        focus: int = 0,
    ):
        # header
        header = urwid.AttrMap(urwid.Text(title, align="center"), "reversed")
        footer = urwid.Text(statusmsg)
        # build the question between header and menu
        body = [
            urwid.Divider(),
            urwid.Text(question),
            urwid.Divider(),
        ]
        # build the menu
        buttons = []
        for c in choices:
            button = urwid.Button(c)
            urwid.connect_signal(button, "click", self.item_chosen, c)
            buttons.append(urwid.AttrMap(button, None, focus_map="reversed"))
        pile = urwid.Pile(buttons)
        pile.set_focus(focus)
        body.append(urwid.LineBox(pile))
        # build the help message between menu and footer
        body.extend(
            [
                urwid.Divider(),
                urwid.Text(helpmsg),
            ]
        )
        # assemble the widgets
        body = ListBoxVimKeys(urwid.SimpleFocusListWalker(body))
        frame = urwid.Frame(header=header, body=body, footer=footer)
        frame = urwid.Padding(frame, left=1, right=1)
        loop = urwid.MainLoop(
            frame,
            palette=[("reversed", "standout", "")],
            unhandled_input=self.exit_on_esc,
        )
        self.loop = loop

    def run(self):
        try:
            self.loop.run()
        except urwid.ExitMainLoop:
            return None
        return self._choice


class CtrlUEdit(urwid.Edit):
    def keypress(self, size, key) -> Optional[str]:
        if key == "ctrl u":
            self.set_edit_text("")
            return None
        return super().keypress(size, key)


class SingleEdit(object):
    # we store the user choice here
    _choice: Optional[str] = None

    @staticmethod
    def exit_on_esc(key: str) -> None:
        if key in ("esc",):
            sys.exit(1)
        elif key in ("enter",):
            raise urwid.ExitMainLoop()

    def edit_update(self, edit: urwid.Edit, new_edit_text: str) -> None:
        self._choice = new_edit_text

    def __init__(
        self, title: str, question: str, default: str, helpmsg: str, statusmsg: str
    ):
        # process arguments
        self._choice = default
        # header and footer
        header = urwid.AttrMap(urwid.Text(title, align="center"), "reversed")
        footer = urwid.Text(statusmsg)
        # build the question between header and menu
        body = [
            urwid.Divider(),
            urwid.Padding(urwid.Text(question), left=1, right=1),
            urwid.Divider(),
        ]
        # build the edit widget
        edit = CtrlUEdit("", default)
        urwid.connect_signal(edit, "change", self.edit_update)
        edit = urwid.LineBox(edit)
        body.append(edit)
        body.extend(
            [
                urwid.Divider(),
                urwid.Padding(urwid.Text(helpmsg), left=1, right=1),
            ]
        )
        # assemble the widgets
        body = urwid.ListBox(urwid.SimpleFocusListWalker(body))
        frame = urwid.Frame(header=header, body=body, footer=footer)
        frame = urwid.Padding(frame, left=1, right=1)
        loop = urwid.MainLoop(
            frame,
            palette=[("reversed", "standout", "")],
            unhandled_input=self.exit_on_esc,
        )
        self.loop = loop

    def run(self):
        self.loop.run()
        return self._choice


def _request_frontend_specific_config(
    frontend: str,
    current_config: Dict = dict(),
) -> dict:
    """
    ask the user to provide the frontend-specific configuration
    """
    conf = dict()

    # openai part
    if frontend == "openai" and "openai_base_url" not in current_config:
        value = SingleEdit(
            _TITLE,
            "Enter the OpenAI base url:",
            default["openai_base_url"],
            "Keep the default as is, if you do not intend to use this API on a different compatible service.",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["openai_base_url"] = value
    if frontend == "openai" and "openai_api_key" not in current_config:
        value = SingleEdit(
            _TITLE,
            "Enter the OpenAI API key:",
            default["openai_api_key"],
            "Typically your key can be found here: https://platform.openai.com/settings/organization/api-keys",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["openai_api_key"] = value
    if frontend == "openai" and "openai_model" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the OpenAI API model name:",
            default["openai_model"],
            "If not sure, just keep the default. Available options: https://platform.openai.com/docs/models",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["openai_model"] = value
    # anthropic part
    if frontend == "anthropic" and "anthropic_api_key" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the Anthropic API key:",
            default["anthropic_api_key"],
            "Typicall your key can be found here: https://console.anthropic.com/settings/keys",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["anthropic_api_key"] = value
    if frontend == "anthropic" and "anthropic_model" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the Anthropic model name:",
            default["anthropic_model"],
            "If not sure, just keep the default. Available options: https://docs.anthropic.com/en/docs/about-claude/models",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["anthropic_model"] = value

    # google part
    if frontend == "google" and "google_api_key" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the Google Gemini API key:",
            default["google_api_key"],
            "Typically found here: https://aistudio.google.com/app/apikey",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["google_api_key"] = value
    if frontend == "google" and "google_model" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the Google model name:",
            default["google_model"],
            "If not sure, just keep the default. Available options: https://ai.google.dev/gemini-api/docs/models/gemini",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["google_model"] = value

    # deepseek part
    if frontend == "deepseek" and "deepseek_base_url" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the DeepSeek service url:",
            default["deepseek_base_url"],
            "Just keep the default if you intend to use DeepSeek API service. Reference: https://api-docs.deepseek.com/",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["deepseek_base_url"] = value

    if frontend == "deepseek" and "deepseek_api_key" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the DeepSeek API key:",
            default["deepseek_api_key"],
            "Typically found here: https://platform.deepseek.com/api_keys",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["deepseek_api_key"] = value

    if frontend == "deepseek" and "deepseek_model" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the DeepSeek model name:",
            default["deepseek_model"],
            "If not sure, just keep the default. As of Jan 24 2025, the available models are: deepseek-chat, deepseek-reasoner. The list of latest models can be found here: https://api-docs.deepseek.com/quick_start/pricing",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["deepseek_model"] = value

    # vLLM (or any other OpenAI-API compatible service)
    if frontend == "vllm" and "vllm_base_url" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the vLLM (or any OpenAI-API compatible) service url:",
            default["vllm_base_url"],
            "Reference: https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["vllm_base_url"] = value
    if frontend == "vllm" and "vllm_api_key" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the vLLM (or any OpenAI-API compatible) API key:",
            default["vllm_api_key"],
            "Reference: https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["vllm_api_key"] = value
    if frontend == "vllm" and "vllm_model" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the vLLM (or any OpenAI-API compatible) model name:",
            default["vllm_model"],
            "Reference: https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["vllm_model"] = value

    # zmq part
    if frontend == "zmq" and "zmq_backend" not in current_config:
        value = SingleEdit(
            "DebGPT Configurator",
            "Enter the DebGPT ZMQ Backend URL:",
            default["zmq_backend"],
            "The service endpoint where you launched debgpt backend.",
            "Press Enter to confirm. Press Esc to abort.",
        ).run()
        _abort_on_None(value)
        conf["zmq_backend"] = value

    # dryrun part
    if frontend == "dryrun":
        pass

    return conf


def _edit_config_template(config_template: str, key: str, value: str) -> str:
    """
    Edit the configuration template (toml file). Updating the value of
    the specified key.

    Args:
        config_template: the configuration template as a string
        key: the key to be updated
        value: the value to be updated
    Returns:
        the updated configuration template
    """
    lines = config_template.split("\n")
    newlines = []
    for line in lines:
        if re.match(r"^\s*{}\s*=".format(key), line):
            line = "{} = {}".format(key, value)
        newlines.append(line)
    return "\n".join(newlines)


def fresh_install_guide(dest: Optional[str] = None) -> dict:
    """
    This function is a configuration guide for fresh installation of DebGPT.
    """
    conf = dict()
    config_template: str = arguments.parse_args([]).config_template

    # step 1: select a frontend
    frontends = [
        "OpenAI    | For GPT models, or any other OpenAI-API & compatible services",
        "Anthropic | For Claude models through Anthropic-API",
        "Google    | For Gemini models through Google-API",
        "DeepSeek  | For DeepSeek models through OpenAI-API",
        "vLLM      | self-hosted, or any other OpenAI-API compatible services",
        "ZMQ       | self-hosted, DebGPT built-in LLM backend",
        "Dryrun    | DebGPT built-in, debugging purpose only",
    ]
    frontends_focus = {
        "openai": 0,
        "anthropic": 1,
        "google": 2,
        "deepseek": 3,
        "vllm": 4,
        "zmq": 5,
        "dryrun": 6,
    }[default["frontend"]]

    frontend = SingleChoice(
        "DebGPT Configurator",
        "Select a frontend that DebGPT will use:",
        frontends,
        "A frontend is a client that communicates with \
its corresponding backend that serves large language model (LLM). \
To use a commercial \
service, you may need to sign up and pay for an API key. Besides, \
if you have a spare GPU or a powerful CPU, you can take a look at the \
self-hosted LLM services. A web search can direct you to the details.\n\n\
OpenAI frontend: besides the official OpenAI service, this frontend works \
with any service that exposes an OpenAI-compatible API endpoint — simply \
provide the appropriate base URL in the next step.\n\n\
This configurator will generate a minimal configuration file for you to \
make DebGPT work with the selected frontend.\n\n\
For advanced usages and more options, you may generate a configuration \
template with the following command for manual editing:\n\n\
  $ debgpt genconfig > ~/.config/debgpt/config.toml\n\n\
This could be useful if you wish to switch among multiple frontends \
using the `--frontend|-F` argument.",
        "Press Enter to confirm. Press Esc to abort.",
        focus=frontends_focus,
    ).run()
    _abort_on_None(frontend)
    frontend = frontend.split(" ")[0].lower().replace(".", "")
    conf["frontend"] = frontend

    # step 2: ask for the frontend-specific configuration
    extra = _request_frontend_specific_config(frontend)
    conf.update(extra)

    # edit the configuration template
    for k, v in conf.items():
        if isinstance(v, bool):
            v = "true" if v else "false"
            config_template = _edit_config_template(config_template, k, v)
        elif isinstance(v, str):
            config_template = _edit_config_template(config_template, k, repr(v))
        else:
            raise ValueError("Unexpected type:", type(v))
    config_template += "\n"  # new line at the end

    # final: write configuration to specified destination
    if dest:
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        with open(dest, "wt") as f:
            f.write(config_template)
        console.log("Config written to:", dest)
        console.print("[white on violet]>_< Enjoy DebGPT!")
    else:
        # verbose print
        console.print("Configuration (config.toml):")
        print("```")
        print(config_template)
        print("```")

    return conf


if __name__ == "__main__":
    miniconfig = fresh_install_guide()
