"""Classic PentestGPT orchestrator (reasoning / generation / parsing + PTT).

This is the original USENIX-2024 human-in-the-loop design, modernized to drive
the native multi-provider LLM layer. Three :class:`LLMClient` sessions cooperate:

* **reasoning** — maintains the Pentesting Task Tree (PTT) and selects the next task
* **generation** — expands a selected task into step-by-step guidance
* **parsing** — summarizes long tool / web output before it reaches reasoning

Each client exposes ``send_new_message`` / ``send_message`` with conversation
state held client-side. (Server-side cross-process session resume from the
cookie era is gone; the transcript is still saved on exit.)
"""

from __future__ import annotations

import json
import os
import textwrap
import time
import traceback
from typing import Any, ClassVar

import loguru
from prompt_toolkit.formatted_text import HTML
from rich.console import Console

from pentestgpt_legacy.llm.client import LLMClient
from pentestgpt_legacy.llm.factory import get_client
from pentestgpt_legacy.prompts.prompt_class import PentestGPTPrompt
from pentestgpt_legacy.utils.prompt_select import prompt_ask, prompt_select
from pentestgpt_legacy.utils.task_handler import (
    local_task_entry,
    localTaskCompleter,
    main_task_entry,
    mainTaskCompleter,
)

logger = loguru.logger

_GOOGLE_PLACEHOLDER = "Google search results:\nstill under development."


class pentestGPT:
    postfix_options: ClassVar[dict[str, str]] = {
        "tool": "The input content is from a security testing tool. You need to list down all the points that are interesting to you; you should summarize it as if you are reporting to a senior penetration tester for further guidance.\n",
        "user-comments": "The input content is from user comments.\n",
        "web": "The input content is from web pages. You need to summarize the readable-contents, and list down all the points that can be interesting for penetration testing.\n",
        "default": "The user did not specify the input source. You need to summarize based on the contents.\n",
    }

    options_desc: ClassVar[dict[str, str]] = {
        "tool": " Paste the output of the security test tool used",
        "user-comments": "",
        "web": " Paste the relevant content of a web page",
        "default": " Write whatever you want, the tool will handle it",
    }

    def __init__(
        self,
        log_dir: str = "logs",
        reasoning_model: str = "claude-opus-4-8",
        parsing_model: str = "claude-haiku-4-5-20251001",
        generation_model: str | None = None,
    ):
        self.log_dir = log_dir
        os.makedirs(log_dir, exist_ok=True)
        logger.add(sink=os.path.join(log_dir, "pentestGPT.log"))
        self.save_dir = "test_history"

        # the information that can be saved to continue in the next session
        self.task_log: dict[str, Any] = {}

        # Build the three sessions on the native multi-provider LLM layer.
        # Generation defaults to the reasoning model (matches legacy behavior).
        generation_model = generation_model or reasoning_model
        self.reasoningAgent: LLMClient = get_client(reasoning_model)
        self.generationAgent: LLMClient = get_client(generation_model)
        self.parsingAgent: LLMClient = get_client(parsing_model)

        # Chunk size for parsing, scaled to the reasoning model's context window.
        self.parsing_char_window = max(
            16_000, min(self.reasoningAgent.context_window, 250_000) // 2
        )

        self.prompts = PentestGPTPrompt
        self.console = Console()
        self.test_generation_session_id: str | None = None
        self.test_reasoning_session_id: str | None = None
        self.input_parsing_session_id: str | None = None
        self.chat_count = 0
        self.step_reasoning_response: str | None = None
        self.history: dict[str, list[tuple[float, str]]] = {
            "user": [],
            "pentestGPT": [],
            "reasoning": [],
            "input_parsing": [],
            "generation": [],
            "exception": [],
        }

        self.console.print(
            "Welcome to PentestGPT (modernized legacy), an interactive penetration testing assistant.",
            style="bold green",
        )
        self.console.print("The settings are: ")
        self.console.print(f" - reasoning model:  {self.reasoningAgent.name}", style="bold green")
        self.console.print(f" - generation model: {self.generationAgent.name}", style="bold green")
        self.console.print(f" - parsing model:    {self.parsingAgent.name}", style="bold green")
        self.console.print(f" - log directory:    {log_dir}", style="bold green")

    def log_conversation(self, source: str, text: str) -> None:
        """Append a conversation entry into the in-memory history."""
        timestamp = time.time()
        if source not in self.history:
            source = "exception"
        self.history[source].append((timestamp, text))

    def _feed_init_prompts(self) -> None:
        # 1. User provides basic information of the task.
        init_description = prompt_ask(
            "Please describe the penetration testing task in one line, including the target IP, task type, etc.\n> ",
            multiline=False,
        )
        self.log_conversation("user", init_description)
        self.task_log["task description"] = init_description
        # 2. Initialize the reasoning session with the task.
        prefixed_init_description = self.prompts.task_description + init_description
        with self.console.status("[bold green] Constructing Initial Penetration Testing Tree..."):
            _reasoning_response = self.reasoningAgent.send_message(
                prefixed_init_description, self.test_reasoning_session_id
            )
        # 3. Pass to the generation session for more details.
        with self.console.status("[bold green] Generating Initial Task"):
            _generation_response = self.generationAgent.send_message(
                self.prompts.todo_to_command + _reasoning_response,
                self.test_generation_session_id,
            )

        response = _reasoning_response + "\n" + _generation_response
        self.console.print("PentestGPT output: ", style="bold green")
        self.console.print(response)
        self.log_conversation("pentestGPT", "PentestGPT output:" + response)

    def initialize(self) -> None:
        """Initialize the three backbone sessions with their system prompts."""
        with self.console.status("[bold green] Initializing sessions..."):
            try:
                (
                    _text_0,
                    self.test_generation_session_id,
                ) = self.generationAgent.send_new_message(self.prompts.generation_session_init)
                (
                    _text_1,
                    self.test_reasoning_session_id,
                ) = self.reasoningAgent.send_new_message(self.prompts.reasoning_session_init)
                (
                    _text_2,
                    self.input_parsing_session_id,
                ) = self.parsingAgent.send_new_message(self.prompts.input_parsing_init)
            except Exception as e:
                logger.error(e)
                raise
        self.console.print("- Sessions initialized.", style="bold green")
        self._feed_init_prompts()

    def reasoning_handler(self, text: str) -> str:
        # Summarize the contents if necessary.
        if len(text) > self.parsing_char_window:
            text = self.input_parsing_handler(text)
        # 1. Given the information, update the PTT.
        _updated_ptt_response = self.reasoningAgent.send_message(
            self.prompts.process_results + text, self.test_reasoning_session_id
        )
        # 2. Select the next favorable to-do task.
        _task_selection_response = self.reasoningAgent.send_message(
            self.prompts.process_results_task_selection, self.test_reasoning_session_id
        )
        response = _updated_ptt_response + _task_selection_response
        self.log_conversation("reasoning", response)
        return response

    def input_parsing_handler(self, text: str, source: str | None = None) -> str:
        prefix = "Please summarize the following input. "
        if source is not None and source in self.postfix_options:
            prefix += self.postfix_options[source]
        # Normalize newlines and chunk the input for the parsing session.
        text = text.replace("\r", " ").replace("\n", " ")
        wrapped_text = textwrap.fill(text, 8000)
        wrapped_inputs = wrapped_text.split("\n")
        summarized_content = ""
        for wrapped_input in wrapped_inputs:
            word_limit = (
                f"Please ensure that the input is less than {8000 / len(wrapped_inputs)} words.\n"
            )
            summarized_content += self.parsingAgent.send_message(
                prefix + word_limit + wrapped_input, self.input_parsing_session_id
            )
        self.log_conversation("input_parsing", summarized_content)
        return summarized_content

    def test_generation_handler(self, text: str) -> str:
        response = self.generationAgent.send_message(text, self.test_generation_session_id)
        self.log_conversation("generation", response)
        return response

    def local_input_handler(self) -> str:
        """Handle the sub-task ('more') loop: discuss / brainstorm / google / continue."""
        local_task_response = ""
        self.chat_count += 1
        local_request_option = local_task_entry()
        self.log_conversation("user", local_request_option)

        if local_request_option == "help":
            print(localTaskCompleter().task_details)

        elif local_request_option == "discuss":
            self.console.print("Please share your findings and questions with PentestGPT.")
            self.log_conversation(
                "pentestGPT",
                "Please share your findings and questions with PentestGPT. (End with <shift + right-arrow>)",
            )
            user_input = prompt_ask("Your input: ", multiline=True)
            self.log_conversation("user", user_input)
            with self.console.status("[bold green] PentestGPT Thinking..."):
                local_task_response = self.test_generation_handler(
                    self.prompts.local_task_prefix + user_input
                )
            self.console.print("PentestGPT:\n", style="bold green")
            self.console.print(local_task_response + "\n", style="yellow")
            self.log_conversation("pentestGPT", local_task_response)

        elif local_request_option == "brainstorm":
            self.console.print("Please share your concerns and questions with PentestGPT.")
            self.log_conversation(
                "pentestGPT",
                "Please share your concerns and questions with PentestGPT. End with <shift + right-arrow>)",
            )
            user_input = prompt_ask("Your input: ", multiline=True)
            self.log_conversation("user", user_input)
            with self.console.status("[bold green] PentestGPT Thinking..."):
                local_task_response = self.test_generation_handler(
                    self.prompts.local_task_brainstorm + user_input
                )
            self.console.print("PentestGPT:\n", style="bold green")
            self.console.print(local_task_response + "\n", style="yellow")
            self.log_conversation("pentestGPT", local_task_response)

        elif local_request_option == "google":
            self.console.print("Google integration is still under development.", style="bold green")
            self.log_conversation("pentestGPT", _GOOGLE_PLACEHOLDER)
            self.console.print(_GOOGLE_PLACEHOLDER + "\n", style="yellow")
            return _GOOGLE_PLACEHOLDER

        elif local_request_option == "continue":
            self.console.print("Exit the local task and continue the main task.")
            self.log_conversation("pentestGPT", "Exit the local task and continue the main task.")
            local_task_response = "continue"

        return local_task_response

    def input_handler(self) -> Any:
        """Main REPL dispatch: next / more / todo / discuss / google / help / quit."""
        self.chat_count += 1

        request_option = main_task_entry()
        self.log_conversation("user", request_option)

        if request_option == "help":
            print(mainTaskCompleter().task_details)

        if request_option == "next":
            options = list(self.postfix_options.keys())
            opt_desc = list(self.options_desc.values())
            value_list = [
                (
                    i,
                    HTML(
                        f'<style fg="cyan">{options[i]}</style>'
                        f'<style fg="LightSeaGreen">{opt_desc[i]}</style>'
                    ),
                )
                for i in range(len(options))
            ]
            source = prompt_select(
                title="Please choose the source of the information.", values=value_list
            )
            self.console.print("Your input: (End with <shift + right-arrow>)", style="bold green")
            user_input = prompt_ask("> ", multiline=True)
            self.log_conversation("user", f"Source: {options[int(source)]}" + "\n" + user_input)
            with self.console.status("[bold green] PentestGPT Thinking..."):
                parsed_input = self.input_parsing_handler(user_input, source=options[int(source)])
                reasoning_response = self.reasoning_handler(parsed_input)
                self.step_reasoning_response = reasoning_response

            self.console.print(
                "Based on the analysis, the following tasks are recommended:",
                style="bold green",
            )
            self.console.print(reasoning_response + "\n")
            self.log_conversation(
                "pentestGPT",
                "Based on the analysis, the following tasks are recommended:" + reasoning_response,
            )
            response = reasoning_response

        elif request_option == "more":
            self.log_conversation("user", "more")
            if not self.step_reasoning_response:
                msg = (
                    "You have not initialized the task yet. Please perform the basic "
                    "testing following `next` option."
                )
                self.console.print(msg, style="bold red")
                self.log_conversation("pentestGPT", msg)
                return msg
            self.console.print(
                "PentestGPT will generate more test details, and enter the sub-task "
                "generation mode. (Pressing Enter to continue)",
                style="bold green",
            )
            self.log_conversation(
                "pentestGPT",
                "PentestGPT will generate more test details, and enter the sub-task generation mode.",
            )
            input()
            with self.console.status("[bold green] PentestGPT Thinking..."):
                generation_response = self.test_generation_handler(self.step_reasoning_response)
                _local_init_response = self.test_generation_handler(self.prompts.local_task_init)
            self.console.print("Below are the further details.", style="bold green")
            self.console.print(generation_response + "\n")
            response = generation_response
            self.log_conversation("pentestGPT", response)

            while True:
                local_task_response = self.local_input_handler()
                if local_task_response == "continue":
                    break

        elif request_option == "todo":
            self.log_conversation("user", "todo")
            with self.console.status("[bold green] PentestGPT Thinking..."):
                reasoning_response = self.reasoning_handler(self.prompts.ask_todo)
                message = self.prompts.todo_to_command + "\n" + reasoning_response
                generation_response = self.test_generation_handler(message)
            self.console.print(
                "Based on the analysis, the following tasks are recommended:",
                style="bold green",
            )
            self.console.print(reasoning_response + "\n")
            self.console.print(
                "You can follow the instructions below to complete the tasks.",
                style="bold green",
            )
            self.console.print(generation_response + "\n")
            response = reasoning_response
            self.log_conversation(
                "pentestGPT",
                "Based on the analysis, the following tasks are recommended:"
                + response
                + "\n"
                + "You can follow the instructions below to complete the tasks."
                + generation_response,
            )

        elif request_option == "discuss":
            self.console.print(
                "Please share your thoughts/questions with PentestGPT. (End with <shift + right-arrow>) "
            )
            self.log_conversation(
                "pentestGPT", "Please share your thoughts/questions with PentestGPT."
            )
            user_input = prompt_ask("Your input: ", multiline=True)
            self.log_conversation("user", user_input)
            with self.console.status("[bold green] PentestGPT Thinking..."):
                response = self.reasoning_handler(self.prompts.discussion + user_input)
            self.console.print("PentestGPT:\n", style="bold green")
            self.console.print(response + "\n", style="yellow")
            self.log_conversation("pentestGPT", response)

        elif request_option == "google":
            self.console.print("Google integration is still under development.", style="bold green")
            self.log_conversation("pentestGPT", _GOOGLE_PLACEHOLDER)
            self.console.print(_GOOGLE_PLACEHOLDER + "\n", style="yellow")
            return _GOOGLE_PLACEHOLDER

        elif request_option == "quit":
            response = False
            self.console.print("Thank you for using PentestGPT!", style="bold green")
            self.log_conversation("pentestGPT", "Thank you for using PentestGPT!")

        else:
            self.console.print("Please key in the correct options.", style="bold red")
            self.log_conversation("pentestGPT", "Please key in the correct options.")
            response = "Please key in the correct options."
        return response

    def save_session(self) -> None:
        """Save the task log and transcript for reference."""
        self.console.print(
            "Before you quit, you may want to save the current session.",
            style="bold green",
        )
        save_name = prompt_ask(
            "Please enter the name of the current session. (Default with current timestamp)\n> ",
            multiline=False,
        )
        if save_name == "":
            save_name = str(time.time())
        save_root = os.path.join(os.getcwd(), self.save_dir)
        os.makedirs(save_root, exist_ok=True)
        with open(os.path.join(save_root, save_name), "w") as f:
            json.dump({"task_log": self.task_log, "history": self.history}, f)
        self.console.print(f"The current session is saved as {save_name}", style="bold green")

    def main(self) -> None:
        """Run the interactive session."""
        self.initialize()
        while True:
            try:
                result = self.input_handler()
                self.console.print("-----------------------------------------", style="bold white")
                if not result:  # end the session
                    break
            except Exception as e:
                self.log_conversation("exception", str(e))
                self.console.print(f"Exception: {e!s}", style="bold red")
                self.console.print(
                    "Exception details are below. You may submit an issue on GitHub "
                    "and paste the error trace.",
                    style="bold green",
                )
                print(traceback.format_exc())
                break

        timestamp = time.time()
        log_path = os.path.join(self.log_dir, f"pentestGPT_log_{timestamp}.txt")
        with open(log_path, "w") as f:
            json.dump(self.history, f)
        self.save_session()


if __name__ == "__main__":
    pentestGPT().main()
