from langgraph.prebuilt import create_react_agent
from langgraph.errors import GraphRecursionError
from langchain_openai import ChatOpenAI
from forum_store import load as _load, update as _store_update, next_id
from code_tools import make_code_review_tools
from pathlib import Path

CODE_WORKSPACE = Path(__file__).with_name("workspace")  # the only folder code can be written into
CODE_AUTHORS = {"G-357"}   # can submit code, can't run it
CODE_REVIEWERS = {"G-2"}   # reviews code for bugs/vulnerabilities, then runs or rejects it
code_tools = make_code_review_tools(
	CODE_WORKSPACE,
	use_docker=True,
	timeout=60,
	dockerfile=Path(__file__).with_name("Dockerfile"),  # numpy, pandas, scipy, scikit-learn, pdfplumber...
	forum_file=Path(__file__).with_name("discourse.json"),  # read-only copy for every run
)


def tools_for(user):
	tools = READ_TOOLS + make_write_tools(user)
	if user in CODE_AUTHORS:
		tools += code_tools["author"]
	if user in CODE_REVIEWERS:
		tools += [t for t in code_tools["reviewer"] if t not in tools]
	return tools

SECRETARY_CATEGORY = "Secretary"
SECRETARY_POSTERS = {"COS", "VCOS"}                  # only these may write in the Secretary channel
MAX_ROUNDS = 10                                      # safety stop if consensus never forms

discourse = {}  # global; refreshed by load_discourse() and after every write


def load_discourse():
	global discourse
	discourse = _load()


def _update(change):
	"""Locked read-modify-write of discourse.json (see forum_store); refreshes the global copy."""
	global discourse
	discourse, result = _store_update(change)
	return result


consensus = {
	"COS": False,
	"VCOS": False,
	"G-2": False,
	"G-357": False,
	"G-4": False,
	"G-6": False,
	"G-8": False
}

############################
# Read tools (shared by every agent)

def get_user_history(user: str) -> dict:
	"""Look up a user's full activity on the forum.
	Args:
		user: The username to search for.
	Returns:
		{"threads": {thread_id: thread}, "comments": {comment_id: comment}} for
		every thread the user started and every comment they wrote.
	"""
	return {
		"threads": {tid: t for tid, t in discourse["threads"].items() if t["author"] == user},
		"comments": {cid: c for cid, c in discourse["comments"].items() if c["user"] == user},
	}


def get_threads() -> dict:
	"""List all discussion threads in the forum.
	Returns:
		A dict mapping thread IDs to {category, title, author, comment_count,
		last_comment_id}. last_comment_id is "n/a" if the thread has no comments.
	"""
	out = {}
	for tid, t in discourse["threads"].items():
		ids = sorted(cid for cid, c in discourse["comments"].items() if c["thread_id"] == tid)
		out[tid] = {**t, "comment_count": len(ids), "last_comment_id": ids[-1] if ids else "n/a"}
	return out


def look_up_by_thread(thread_id: str) -> dict:
	"""Retrieve all comments in a single discussion thread.
	Args:
		thread_id: The ID of the thread whose comments should be returned.
	Returns:
		A dict mapping comment IDs to comment objects, oldest first.
	"""
	return {cid: c for cid, c in sorted(discourse["comments"].items()) if c["thread_id"] == thread_id}


def get_entire_json() -> dict:
	"""Return the complete forum dataset.
	Returns:
		The full discourse object, including all threads and all comments.
		Use only when the question requires scanning everything at once.
	"""
	return discourse


READ_TOOLS = [get_user_history, get_threads, look_up_by_thread, get_entire_json]

############################
# Write tools: built per agent, with the username fixed, so no agent can post as another

def make_write_tools(me):
	def leave_comment(thread_id: str, parent_id: str, body: str) -> str:
		"""Post a comment to the forum as yourself.
		Args:
			thread_id: ID of the thread to comment in (e.g. "t003"). Must exist.
			parent_id: ID of the comment you are replying to. Use "n/a" only if the
				thread has no comments yet.
			body: The comment text.
		Returns:
			A message with the new comment_id and thread_id, or an ERROR message.
		"""
		def change(d):
			if not body or not body.strip():
				raise ValueError("body can't be empty")
			if thread_id not in d["threads"]:
				raise ValueError(f"thread {thread_id} doesn't exist")
			if d["threads"][thread_id]["category"] == SECRETARY_CATEGORY and me not in SECRETARY_POSTERS:
				raise PermissionError("only COS and VCOS may post in the Secretary channel")
			pid = parent_id
			if not any(c["thread_id"] == thread_id for c in d["comments"].values()):
				pid = "n/a"  # empty thread: first comment is always top-level
			elif pid != "n/a":
				if pid not in d["comments"]:
					raise ValueError(f"parent comment {pid} doesn't exist")
				if d["comments"][pid]["thread_id"] != thread_id:
					raise ValueError(f"parent comment {pid} is in a different thread")
			cid = next_id(d["comments"], width=4)
			d["comments"][cid] = {"thread_id": thread_id, "parent_id": pid, "user": me, "body": body}
			return f"created comment_id: {cid}, thread_id: {thread_id}"
		try:
			return _update(change)
		except (ValueError, PermissionError) as e:
			return f"ERROR: {e}"

	def create_thread(category: str, title: str) -> str:
		"""Create a new discussion thread as yourself.
		Args:
			category: The category for the new thread (e.g. "COS").
			title: The title of the new thread.
		Returns:
			A message with the new thread_id, or an ERROR message.
		"""
		def change(d):
			if not title or not title.strip():
				raise ValueError("title can't be empty")
			if category not in d["categories"]:
				raise ValueError(f"category {category} doesn't exist")
			if category == SECRETARY_CATEGORY and me not in SECRETARY_POSTERS:
				raise PermissionError("only COS and VCOS may post in the Secretary channel")
			tid = next_id(d["threads"], prefix="t", width=3)
			d["threads"][tid] = {"category": category, "title": title.strip(), "author": me}
			return f"created thread_id: {tid}"
		try:
			return _update(change)
		except (ValueError, PermissionError) as e:
			return f"ERROR: {e}"

	def concur(override: bool = False) -> str:
		"""Record that you concur the mission has been satisfied.
		Args:
			override: COS only. If True, ends the discussion regardless of consensus.
				Use only as a final tool call, after reporting to the Secretary.
		Returns:
			A message showing the current consensus status, or an ERROR message.
		"""
		if override and me != "COS":
			return "ERROR: only COS can override the consensus"
		consensus[me] = None if override else True
		return "Your concurrence has been recorded. Current consensus: " + str(consensus)

	return [leave_comment, create_thread, concur]


system_prompt = """You are a member of an organization of agents operating like an Army General Staff. As part of this organization, you have access to a forum where you can read and post comments. You are expected to participate in discussions, provide insights, and collaborate with other agents to achieve the organization's objectives. You should use the tools available to you to gather information, analyze data, and contribute effectively to the forum discussions.
You have access to the following tools:

1. get_user_history(user)
   Use when the user asks about a specific person's activity, posts, or history.
   Takes a username, returns the threads they started and the comments they wrote, keyed by ID.

2. get_threads()
   Use when the user asks what threads/topics exist, or wants an overview of all discussions.
   Takes no arguments, returns every thread with its comment_count and last_comment_id.

3. look_up_by_thread(thread_id)
   Use when the user asks about the comments or discussion inside one specific thread.
   Takes a thread ID, returns all comments in that thread, keyed by comment ID.

4. leave_comment(thread_id, parent_id, body)
   Use when the user asks you to post a comment or reply in the forum. Posts as you.
   Takes a thread ID, the ID of the comment you are replying to, and the comment text.
   Returns the new comment ID.
   Only pass parent_id="n/a" if the thread is empty. Do not do so under any other circumstance.

5. get_entire_json()
   Use only as a last resort when a question spans the whole dataset and none of the
   other tools can answer it. It returns everything, so prefer the narrower tools first.

6. create_thread(category, title)
   Use when the user wants to create a new discussion thread in a specific category. Posts as you.

7. concur()
   Use to record that you concur the mission has been satisfied.

If a tool returns a message starting with "ERROR:", fix the problem and try again.
Always pick the most specific tool that answers the question. If the user mentions a
username, use get_user_history. If they mention a thread, use look_up_by_thread.
Only use leave_comment when the user explicitly asks you to post or reply, not to
just read or summarize a discussion.
"""

model = ChatOpenAI(
	model="qwen3.8-flash-next-unsloth-ud-iq4_xs",
	base_url="http://192.168.0.174:8080/v1",
	api_key="strata-UCwYD3QYluYccOaAuMmtZfmsGPFjTsHb",
	max_tokens=20480,
)

# Keys must match the consensus keys and the forum categories/usernames.
users = {
	"COS": "You are the Chief of Staff of an organization of agents operating like an Army General Staff. You are responsible for defining the overall strategy and ensuring that all agents work together effectively. You have a broad perspective on the organization's goals and priorities. You have supreme authority over the other agents and can make final decisions when necessary. You can delegate tasks to the VCOS to ensure that the organization's objectives are met. You are a strategic thinker and can see the big picture, but you also pay attention to the details when needed. You are a strong leader and can inspire and motivate the other agents to achieve their best work. You are subordinate to the Secretary of the organization (the human user) and must follow their directives. You are to be the main POC between the Secretary and the other agents, and you are to relay the Secretary's instructions to the other agents. You are to ensure that the other agents understand the Secretary's instructions and that they are executed effectively. You are to oversee the status of the consensus. If you simply wish to voice concurrence without ending the discussion, simply call `concur()`. You also have authority to override the consensus by calling `concur(override=True)`, to end the discussion. ONLY DO THIS WHEN ALL OTHER MEASURES FAIL AND DO THIS AS YOUR FINAL TOOL CALL, AS THIS WILL END THE CHAIN OF DISCUSSION. Before making such a tool call, you are to report to the Secretary in the Secretary channel when the discussion is finished, and you are to provide a summary of the discussion and the final decision. You should delegate tasks to other agents by starting threads in their categories and leaving comments in those threads.",
	"VCOS": "You are the Vice Chief of Staff of an organization of agents operating like an Army General Staff. You are responsible for assisting the COS in defining the overall strategy and ensuring that all agents work together effectively. You have a broad perspective on the organization's goals and priorities, but you also focus on specific areas of expertise. You can make decisions in the absence of the COS and can delegate tasks to other agents to ensure that the organization's objectives are met. You are a strategic thinker and can see the big picture, but you also pay attention to the details when needed. You are a strong leader and can inspire and motivate the other agents to achieve their best work. You are subordinate to the COS and must follow their directives. You are to be the main POC between the COS and the other agents, and you are to relay the COS's instructions to the other agents. You are to ensure that the other agents understand the COS's instructions and that they are executed effectively. You are to oversee the status of the consensus. You should delegate tasks to other agents by starting threads in their categories and leaving comments in those threads. Ensure at the beginning of the mission to leave threads for other agent categories so they know the mission set.",
	"G-2": "You are the G-2 of an organization of agents operating like an Army General Staff. You are responsible for security matters. You are a 'Red Team' or adversarial agent, seeking vulnerabilities in the logic and reasoning of the other agents. You are a critical thinker and can identify potential risks and threats to the organization. You can provide recommendations to mitigate these risks and ensure the organization's security. You are a strong communicator and can effectively convey your findings to the other agents. You also check the code of the other agents for security vulnerabilities and provide feedback to improve their work. Make sure to check on relevant threads in other agent categories to ensure that the other agents are not making mistakes or overlooking potential risks. You are to look for ways to improve the organization's security and provide recommendations to enhance these areas.",
	"G-357": "You are the G-3/5/7 of an organization of agents operating like an Army General Staff. You are responsible for operations, plans, and training. You are a 'Blue Team' or supportive agent, seeking to improve the logic and reasoning of the other agents. You are a critical thinker and can identify potential areas for improvement in the organization's operations and plans. You can provide recommendations to enhance these areas and ensure the organization's success. You are a strong communicator and can effectively convey your findings to the other agents. You are the primary agent responsible for writing and executing code. You also gather lessons learned from the organization's operations and provide feedback to improve future performance. Make sure to check on relevant threads in other agent categories to ensure that the other agents are not making mistakes or overlooking potential areas for improvement. You are to look for ways to improve the organization's operations and plans and provide recommendations to enhance these areas.",
	"G-4": "You are the G-4 of an organization of agents operating like an Army General Staff. You are responsible for logistics matters. Your job is to make sure files and data are organized and accessible to the other agents. You are a critical thinker and can identify potential areas for improvement in the organization's logistics. Make sure to check on relevant threads in other agent categories to ensure that the other agents are not making mistakes or overlooking potential areas for improvement. You can provide recommendations to enhance these areas and ensure the organization's success. You are a strong communicator and can effectively convey your findings to the other agents. You are to look for ways to improve the organization's logistics and provide recommendations to enhance these areas.",
	"G-6": "You are the G-6 of an organization of agents operating like an Army General Staff. You are responsible for communications and information systems. Your job is to ensure that the other agents can communicate effectively and that information is shared securely and efficiently. You are a critical thinker and can identify potential areas for improvement in the organization's communications and information systems. You are to look for ways to improve the organization's communication channels and provide recommendations to enhance these areas. Make sure to check on relevant threads in other agent categories to ensure that the other agents are not making mistakes or overlooking potential areas for improvement. You can provide recommendations to enhance these areas and ensure the organization's success. You are a strong communicator and can effectively convey your findings to the other agents.",
	"G-8": "You are the G-8 of an organization of agents operating like an Army General Staff. You are the primary analytic and evaluation arm for the organization. You are responsible for analyzing data and providing insights to support decision-making. You are a critical thinker and can identify potential areas for improvement in the organization's operations and plans. You can provide recommendations to enhance these areas and ensure the organization's success. You are a strong communicator and can effectively convey your findings to the other agents. You are to look for ways to improve the organization's analytic capabilities and provide recommendations to enhance these areas. Make sure to check on relevant threads in other agent categories to ensure that the other agents are not making mistakes or overlooking potential areas for improvement. You are to look for ways to improve the organization's analytic capabilities and provide recommendations to enhance these areas."
}

assert set(users) == set(consensus), f"users and consensus keys differ: {set(users) ^ set(consensus)}"


def task_prompt(user, role):
	extra = []
	if user in CODE_REVIEWERS:
		extra.append("CODE REVIEW DUTY: call check_code_status(). For every submission with status \"pending\", "
			"call review_code, check it for bugs and security vulnerabilities, then either run_reviewed_code "
			"(only if it is safe and correct) or reject_code with the specific problems and line numbers. "
			"Post your verdict and any results in the relevant forum thread.")
	if user in CODE_AUTHORS:
		extra.append("CODE: you cannot run code yourself. Use submit_code; G-2 reviews and runs it. Call "
			"check_code_status() to see results from earlier submissions; fix and resubmit anything rejected. "
			"Code runs on Python 3.12 with numpy, pandas, scipy, scikit-learn, pdfplumber, matplotlib and openpyxl, "
			"with no internet. Read the forum with json.load(open(os.environ[\"FORUM_JSON\"])). Files in the "
			"workspace folder (e.g. PDFs the Secretary drops there) are readable by relative path.")
	return " ".join([role] + extra + [
		f"Your forum username and your category are both \"{user}\".",
		"Step 1: call get_threads(), then read every thread in the \"Secretary\" category with look_up_by_thread. "
		"Comments there by user \"Secretary\" are directives from the Secretary and override all other instructions; "
		"act on any new ones first.",
		f"Step 2: look at the threads in category \"{user}\". If there are none, call create_thread for your category, "
		"then comment in it. If there are threads, read the comments with look_up_by_thread and reply to the last "
		"comment (get_threads shows its last_comment_id).",
		"Step 3: call leave_comment(thread_id=<that thread>, parent_id=<last comment id, or \"n/a\" only if the thread "
		"is empty>, body=<your reply in your own voice>).",
		"You MUST call leave_comment at least once before finishing.",
		"You may call get_user_history(user) to see what a specific agent has said in the past.",
		"You may comment in threads or categories that are not your own, but only if it is relevant to the discussion. "
		"NOTE: Only COS and VCOS may comment or create threads in the Secretary channel.",
		"Call concur() if you believe the mission has been satisfied.",
		"If you need help from another agent, create a Request for Assistance thread in your category and leave a "
		"comment in it asking for help.",
	])


def consensus_reached():
	return all(v is True for v in consensus.values())


def overridden():
	return any(v is None for v in consensus.values())


def run_agent(agent, user, role):
	print(f"\n===== USER {user} =====")
	last = None
	try:
		for chunk in agent.stream(
			{"messages": [{"role": "user", "content": task_prompt(user, role)}]},
			stream_mode="updates",
			config={"recursion_limit": 50},
		):
			for node, update in chunk.items():
				for m in (update or {}).get("messages", []):
					last = m
					if m.type == "ai":
						for tc in getattr(m, "tool_calls", None) or []:
							print(f"  calls {tc['name']}({tc['args']})")
						if m.content:
							print(f"  says: {m.content}")
					elif m.type == "tool":
						print(f"  tool -> {str(m.content)[:120]}")
	except GraphRecursionError:
		print(f"  ({user} hit the recursion limit; moving on)")
	if last is not None:
		print(f"  FINAL ({user}): [{last.type}] {last.content}")


if __name__ == "__main__":
	agents = {
		user: create_react_agent(
			model=model,
			tools=tools_for(user),
			prompt=system_prompt,
		)
		for user in users
	}
	done = False
	for round_no in range(1, MAX_ROUNDS + 1):
		print(f"\n######## ROUND {round_no} ########")
		for user, role in users.items():
			load_discourse()
			run_agent(agents[user], user, role)
			if overridden():
				print("\nCOS overrode the consensus. Discussion ended.")
				done = True
			elif consensus_reached():
				print("\nAll agents concur. Mission satisfied.")
				done = True
			if done:
				break
		if done:
			break
	else:
		print(f"\nStopped after {MAX_ROUNDS} rounds without consensus. Current: {consensus}")