"""
BaselineAutoGenOrchestrator
Standard AutoGen GroupChat orchestration WITHOUT AACP.

This shows what AutoGen coordination looks like by default:
agents exchange full conversational messages in a GroupChat.
Every coordination hop is a full LLM conversation turn --
preamble, content, and often clarifying questions.

This is the "before" state used by comparison.py.
"""

import os
import json
import time
from pathlib import Path

MODEL_COSTS = {
    "gpt-4o-mini":   0.15,
    "gpt-4o":        5.00,
    "gpt-4.1-mini":  0.40,
    "gpt-4.1":       2.00,
}


class BaselineResult:
    def __init__(self, workflow, model):
        self.workflow   = workflow
        self.model      = model
        self.coord_hops = []
        self.agent_hops = []
        self.outputs    = {}

    @property
    def coordination_cost(self):
        return round(sum(h.get("cost_usd", 0) for h in self.coord_hops), 6)

    @property
    def agent_cost(self):
        return round(sum(h.get("cost_usd", 0) for h in self.agent_hops), 6)

    @property
    def total_cost(self):
        return round(self.coordination_cost + self.agent_cost, 6)

    @property
    def coordination_tokens(self):
        return sum(h.get("tokens_in", 0) for h in self.coord_hops)


class BaselineAutoGenOrchestrator:
    """
    Standard AutoGen GroupChat WITHOUT AACP.
    The orchestrator agent writes natural language messages
    to coordinate specialist agents. Every hop is a full
    conversational LLM turn.
    """

    COORD_PROMPTS = {
        "pay_emp":     "Please retrieve all active employee salary records for {period}. Include employee ID, name, department, cost centre, base salary, any delta, and pension rate. Return as JSON.",
        "pay_budget":  "Please retrieve cost centre budget allocations for {period}. Calculate YTD utilisation and flag any cost centres approaching or exceeding their approved budget. Return as JSON.",
        "pay_merge":   "Using the employee and budget data I am providing, calculate the full payroll for {period}. Apply PAYE at 20%, pension deductions, and compute net pay. Flag budget breaches. Return as JSON.",
        "pay_report":  "Please generate an executive payroll summary report for {period} from the payroll data provided. Include key figures, anomalies, and recommended actions. Return as JSON.",
        "pay_log":     "Please write an audit record confirming the payroll run for {period} is complete and ready for review.",
    }

    def __init__(self, model="gpt-4o-mini", api_key=None,
                 data_dir="data", output_dir="output", verbose=False):
        self.model      = model
        self.api_key    = api_key or os.environ.get("OPENAI_API_KEY")
        self.data_dir   = Path(data_dir)
        self.output_dir = Path(output_dir)
        self.verbose    = verbose
        self._cpm       = MODEL_COSTS.get(model, 1.0)
        self.result     = None

    def _make_agent(self, name, system_message):
        import autogen
        return autogen.ConversableAgent(
            name=name,
            system_message=system_message + "\nRespond with valid JSON only. No markdown fences.",
            llm_config={
                "config_list": [{"model": self.model, "api_key": self.api_key}],
                "temperature": 0.3,
                "max_tokens":  1500,
            },
            human_input_mode="NEVER",
            max_consecutive_auto_reply=1,
        )

    def _coord_and_call(self, orchestrator, agent, prompt_key,
                        period, data=None):
        """
        Orchestrator writes a natural language message to the agent.
        This is the coordination hop -- one full LLM conversation turn.
        """
        import autogen, re

        prompt = self.COORD_PROMPTS[prompt_key].format(period=period)
        if data:
            prompt += f"\n\nData provided:\n{json.dumps(data, indent=2)}"

        # Count coordination tokens
        tokens_est = len(prompt.split()) * 2
        cost_est   = tokens_est / 1_000_000 * self._cpm

        self.result.coord_hops.append({
            "prompt_key": prompt_key,
            "tokens_in":  tokens_est,
            "cost_usd":   round(cost_est, 6),
        })

        if self.verbose:
            print(f'  [NL coord] {prompt_key}: "{prompt[:65]}"  ${cost_est:.4f}')

        # Agent call
        proxy = autogen.ConversableAgent(
            name="PROXY",
            llm_config=False,
            human_input_mode="NEVER",
        )

        start = time.time()
        proxy.initiate_chat(agent, message=prompt, max_turns=1, silent=True)
        latency = (time.time() - start) * 1000

        messages = agent.chat_messages.get(proxy, [])
        raw = ""
        for msg in reversed(messages):
            if msg.get("role") == "assistant" and msg.get("content"):
                raw = msg["content"]
                break

        raw = re.sub(r"^```(?:json)?\s*", "", raw.strip())
        raw = re.sub(r"\s*```$", "", raw)

        agent_tokens = len(raw.split()) * 2
        agent_cost   = agent_tokens / 1_000_000 * self._cpm

        self.result.agent_hops.append({
            "prompt_key": prompt_key,
            "tokens_in":  agent_tokens,
            "cost_usd":   round(agent_cost, 6),
            "latency_ms": round(latency, 0),
        })

        try:
            return json.loads(raw.strip())
        except Exception:
            return {"raw": raw[:200]}

    def run_payroll(self, period="2026-03"):
        import csv
        import autogen

        self.result = BaselineResult("payroll", self.model)

        emp_file = self.data_dir / f"employees_{period}.csv"
        bud_file = self.data_dir / f"budgets_{period}.csv"
        emp_data = list(csv.DictReader(open(emp_file))) if emp_file.exists() else []
        bud_data = list(csv.DictReader(open(bud_file))) if bud_file.exists() else []

        hr_agent  = self._make_agent("HR-AGENT",
            "You are an HR payroll specialist. Process payroll data accurately.")
        fin_agent = self._make_agent("FINANCE-AGENT",
            "You are a finance controller. Manage budget data.")

        orchestrator = autogen.ConversableAgent(
            name="ORCHESTRATOR",
            system_message="You coordinate a payroll workflow.",
            llm_config={
                "config_list": [{"model": self.model, "api_key": self.api_key}],
                "temperature": 0.3, "max_tokens": 200,
            },
            human_input_mode="NEVER",
        )

        print(f"\n  --- Payroll workflow (AutoGen GroupChat, natural language) ---")

        r1 = self._coord_and_call(orchestrator, hr_agent,
            "pay_emp", period, {"employees": emp_data})
        r2 = self._coord_and_call(orchestrator, fin_agent,
            "pay_budget", period, {"budgets": bud_data})
        r3 = self._coord_and_call(orchestrator, hr_agent,
            "pay_merge", period, {"employees": r1, "budgets": r2})
        r4 = self._coord_and_call(orchestrator, hr_agent,
            "pay_report", period, {"payroll": r3})

        self.result.outputs = {
            "employees": r1, "budgets": r2,
            "payroll": r3, "report": r4,
        }
        return self.result
