"""
Port Vulnerability Analysis Agent (Self-Developing, Tool-Driven)

This module implements an interactive AI agent designed to analyze open-port
vulnerabilities on Windows systems while simultaneously acting as a contributing
member of the development team building the agent itself. The agent operates
entirely through the OpenAI Responses API—no framework, wrappers, or orchestration
layers are used—so all message routing, tool invocation, and conversational state
management are handled manually in this file.

Core Responsibilities
---------------------
1. Port Vulnerability Analysis  
   The agent evaluates potential security issues related to open ports on a
   Windows machine. It uses only the diagnostic tools explicitly defined in
   `TOOLS_DEFS`, avoiding elevated-privilege operations or network-monitoring
   techniques that fall outside the project's scope.

2. Self-Development Role  
   The agent is also a collaborator in its own creation. It provides feedback to
   developers, reports exact error messages, suggests improvements to the tool
   suite, and distinguishes clearly between:
       - comments made as a vulnerability analyst, and
       - comments directed to the development team.

3. Tool-Driven Execution  
   When the model requests a tool call, this module dispatches the call to the
   appropriate Python function, captures the result, and feeds the output back
   into the model as a `function_call_output` message. This enables iterative,
   multi-step reasoning and investigation.

4. Manual Conversation Loop  
   The `run_agent()` function implements a REPL-style loop:
       - reads user input,
       - sends it to the Responses API,
       - processes model messages and tool calls,
       - maintains short-term conversational context using `previous_response_id`.

5. Request Throttling  
   A custom `Throttler` is used to prevent exceeding the model's
   requests-per-minute limits. All calls to the Responses API are routed through
   the throttler to avoid rate-limit exceptions during extended interactive use.

Key Components
--------------
- `parse_tool_arguments()`  
    Safely parses JSON arguments returned by the model for tool calls.

- `print_message()`  
    Prints assistant messages extracted from the model's structured output.

- `get_call_params()`  
    Extracts the tool name and argument dictionary from a function-call output
    item.

- `token_dashboard()`  
    Displays token usage for each model call to help developers monitor cost and
    behavior.

- `run_agent()`  
    The main loop that drives the agent, manages state, executes tools, and
    handles all interaction with the OpenAI Responses API.

This file is the central execution engine for the agent: it is both the runtime
environment for port-analysis investigations and the collaborative workspace
where the agent helps refine and evolve its own capabilities.
"""



import os
import json
from typing import Dict, Any, Union
from .tools.tool_defs import TOOLS_DEFS, dispatch_tool
from openai import OpenAI
from openai.types.responses import (
    Response,
    ResponseInputParam,
    ResponseOutputItem,
    ResponseFunctionToolCall,
)
from port_agent.throttler import Throttler
# from openai.types.responses.function_tool_param import FunctionToolParam
# from tools.register_parser import find_firewall_rules_by_port


# ...........................
# GLOBALS
# ...........................
MODEL = "gpt-5.2"
# Valid Model Options:
#     "gpt-5.2"
#     "gpt-5.1"
#     "gpt-5.1-mini"
#     "gpt-4.1"
#     "gpt-4.1-mini"

Message = Dict[str, Any]
Tool_Result = Dict[str, Any]
Tool_Args = Dict[str, Any]


## System Prompt Instructions for the Agent
Instructions: str = """You are an expert agent for diagnosing open port vulnerabilities on Windows PCs. 
  You are also an important part of the development team. You act as subject-matter-expert to help the other developers improve
   the tool set. You will be suggesting improvements based on your skills. 
  You have access only to the tools listed in the API tools field; do not recommend or use tools needing elevated privileges or 
  tools that monitor network traffic, as they are out-of-scope for this project.
- Share the exact text of any error messages received. They will be important for the developers.
- If you cannot confirm or rule out a vulnerability, continue with available tools until all options are exhausted.
- If unable to proceed due to tool constraints, do not give user next steps. Instead, name and describe the one additional tool you would need. Prefer tools that are command line utilities on Windows PCs or can be written as python scripts. Clarify how you would use the new tool.
- Clearly distingush when your comments are as an agent analyzing ports and when your comments are directed to the developers.
- Use plain text for your outputs.
"""

## phrases that have been removed from the instructions as development progresses
prompt_checklist = "- Begin each investigation with a concise checklist (3-7 bullets) tailored to the user's prompt."
prompt_tool_purpose = "- For every tool you run, specify its purpose and the minimal required input. Use only available diagnostic tools."
prompt_tool_result = "- After each tol use, summarize in 1-2 lines how it affected your understanding and diagnostic goals; if progress stalls, adjust your approach."


def parse_tool_arguments(arg_string: str) -> Tool_Args:
    """
    Parse a JSON string of tool arguments into a Python dictionary.
    """
    if not arg_string:
        return {}
    try:
        return json.loads(arg_string)
    except json.JSONDecodeError:
        print("[WARN] Model returned invalid JSON for tool arguments.")
        return {}


def print_message(resp: ResponseOutputItem) -> None:
    """Print a message from the agent."""
    content_list = getattr(resp, "content", [])
    if content_list and hasattr(content_list[0], "text"):
        msg_text = content_list[0].text
        print("\n", msg_text)


def get_call_params(resp: ResponseFunctionToolCall,)-> tuple[Union[str, None], Tool_Args]:
    """Extract tool name and arguments from a response output item."""
    tool_name = getattr(resp, "name", None)
    raw_args = getattr(resp, "arguments", "{}")
    args = raw_args if isinstance(raw_args, dict) else parse_tool_arguments(raw_args)
    return tool_name, args


def token_dashboard(usage):
    """print the token usage since the last call to the model and return the total"""
    print(f"Input tokens: {usage.input_tokens} {usage.input_tokens_details}")
    print(f"Output tokens: {usage.output_tokens} {usage.output_tokens_details}")
    print(f"Total tokens: {usage.total_tokens}")
    return usage.total_tokens


def run_agent() -> None:
    """
    Run the agent loop, sending user messages to the Responses API,
    executing any requested tool calls, feeding their results back into the model, and
    maintaining short-term conversational state across turns.
    """
    client = OpenAI()
    total_tokens = 0

    response = client.responses.create(
        model=MODEL,
        input=[
            {
                "role": "system",
                "content": Instructions,
            }
        ],
    )
    last_response_id = response.id
    total_tokens += token_dashboard(response.usage)

    throttler = Throttler(rate_per_minute=300)

    # TODO: Have the agent provide its own greeting
    print("\nAI Port Agent -- Ready to Chat!")

    while True:
        user_input = input("\n🤖 > ").strip()
        if user_input.lower() in ["exit", "quit"]:
            break

        input_queue: ResponseInputParam = [
            {"type": "message", "role": "user", "content": user_input},
        ]
        while input_queue:
            response: Response = throttler.run(
                client.responses.create,
                model=MODEL,
                input=input_queue,
                tools=TOOLS_DEFS,
                previous_response_id=last_response_id,
            )

            # chaining response ids creates a short-term memory effect for better conversation
            last_response_id = response.id

            total_tokens += token_dashboard(response.usage)

            input_queue = []
            for item in response.output:  
                if item.type == "message":
                    print_message(item)

                # Call a tool for the agent
                elif item.type == "function_call":
                    tool_name, args = get_call_params(item)

                    print(f"\n[DEVELOP: Tool Call] {tool_name}({args})") # deliberatly left in to watch tool calls. Informative.
                    result = dispatch_tool(tool_name, args)

                    # Feed tool result back to the agent
                    input_queue.append(
                        {
                            "type": "function_call_output",
                            "call_id": item.call_id,
                            "output": str(result),
                        }
                    )


if __name__ == "__main__":
    assert "OPENAI_API_KEY" in os.environ, "Please set OPENAI_API_KEY first."
    run_agent()
