"""
Multi-Agent Autonomous Procurement System
LangGraph + LangChain (Groq / Llama-3.3-70b) + LangSmith
"""

import os
import re
import json
import uuid
from typing import TypedDict, Annotated
from datetime import datetime
from operator import add

# ── LangSmith must be configured BEFORE any LangChain import ──────────────────
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"]    = "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"
os.environ["LANGCHAIN_PROJECT"]    = "multi_agent_procurement"
os.environ["LANGCHAIN_ENDPOINT"]   = "https://api.smith.langchain.com"

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

from langchain_groq import ChatGroq
from langchain_core.language_models import BaseLanguageModel
from langchain_core.messages import HumanMessage

# ── Credentials ────────────────────────────────────────────────────────────────
GROQ_API_KEY = "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"
GROQ_MODEL   = "llama-3.3-70b-versatile"
os.environ["GROQ_API_KEY"] = GROQ_API_KEY


# =============================================================================
# 1. STATE
# =============================================================================

class ProcurementState(TypedDict):
    raw_input:           str
    procurement_request: str
    vendor_options:      list
    selected_vendor:     dict
    budget_limit:        float
    analysis_approved:   bool
    human_approved:      bool
    contract_draft:      str
    logs:                Annotated[list, add]


def _empty_state() -> dict:
    return {
        "raw_input":           "",
        "procurement_request": "",
        "vendor_options":      [],
        "selected_vendor":     {},
        "budget_limit":        5000.0,
        "analysis_approved":   False,
        "human_approved":      False,
        "contract_draft":      "",
        "logs":                [],
    }


# =============================================================================
# 2. VENDOR DATABASE (plain functions — no @tool wrapper)
# =============================================================================

def mock_vendor_search(query: str) -> list:
    db = {
        "cloud_infrastructure": [
            {
                "vendor_id": "VENDOR_001",
                "name": "CloudTech Solutions",
                "service_type": "Cloud Infrastructure",
                "price_per_month": 5000,
                "capabilities": ["Auto-scaling", "99.99% Uptime SLA", "24/7 Support"],
                "reputation_score": 9.2,
                "contract_terms": "12-month minimum",
            },
            {
                "vendor_id": "VENDOR_002",
                "name": "InfraSpeed Inc",
                "service_type": "Cloud Infrastructure",
                "price_per_month": 3500,
                "capabilities": ["Auto-scaling", "99.9% Uptime SLA", "Business hours support"],
                "reputation_score": 8.5,
                "contract_terms": "6-month minimum",
            },
            {
                "vendor_id": "VENDOR_003",
                "name": "EcoCloud Global",
                "service_type": "Cloud Infrastructure",
                "price_per_month": 6200,
                "capabilities": ["Auto-scaling", "99.999% Uptime SLA", "24/7 Premium Support", "Green energy"],
                "reputation_score": 9.5,
                "contract_terms": "12-month minimum",
            },
        ],
        "software_licensing": [
            {
                "vendor_id": "VENDOR_004",
                "name": "ProSoft Licensing",
                "service_type": "Enterprise Software",
                "price_per_month": 2000,
                "capabilities": ["1000 user licenses", "Cloud deployment", "Annual updates"],
                "reputation_score": 8.8,
                "contract_terms": "Annual renewal",
            },
            {
                "vendor_id": "VENDOR_005",
                "name": "NextGen Software Co",
                "service_type": "Enterprise Software",
                "price_per_month": 2800,
                "capabilities": ["2000 user licenses", "Cloud + On-prem", "Quarterly updates", "Custom integrations"],
                "reputation_score": 9.1,
                "contract_terms": "Annual renewal",
            },
        ],
    }
    q = query.lower()
    category = "cloud_infrastructure" if any(w in q for w in ["cloud", "server", "infra", "hosting"]) else "software_licensing"
    return db[category]


def validate_budget(vendor_price: float, budget_limit: float) -> dict:
    variance = budget_limit - vendor_price
    within   = vendor_price <= budget_limit
    pct      = round(variance / budget_limit * 100, 2) if budget_limit > 0 else 0
    return {
        "within_budget":       within,
        "vendor_price":        vendor_price,
        "budget_limit":        budget_limit,
        "variance":            variance,
        "variance_percentage": pct,
        "status":              "APPROVED" if within else "REJECTED",
    }


# =============================================================================
# 3. LLM HELPER — plain text call, strips markdown fences, returns str
# =============================================================================

def _llm_call(llm: ChatGroq, prompt: str) -> str:
    response = llm.invoke([HumanMessage(content=prompt)])
    text = response.content.strip()
    # Strip ```json ... ``` or ``` ... ```
    text = re.sub(r"^```[a-zA-Z]*\n?", "", text)
    text = re.sub(r"\n?```$", "", text)
    return text.strip()


# =============================================================================
# 4. AGENT NODES
# =============================================================================

class ProcurementAgents:
    def __init__(self, llm: ChatGroq):
        self.llm = llm

    # ── Agent 0: Prose → Structured ──────────────────────────────────────────
    def parse_node(self, state: ProcurementState) -> Command:
        raw = (state.get("raw_input") or "").strip()
        req = (state.get("procurement_request") or "").strip()

        ts  = datetime.now().strftime("%H:%M:%S")

        # If no prose, just pass through whatever structured request was given
        if not raw:
            log = f"[{ts}] ParseNode: Structured input — skipping prose extraction."
            return Command(
                update={"logs": [log]},
                goto="research_node",
            )

        log = f"[{ts}] ParseNode: Extracting structured fields from prose..."

        prompt = (
            "You are a procurement intake assistant. Extract structured procurement details "
            "from the user description below.\n\n"
            f"User input: {raw}\n\n"
            "Reply with ONLY valid JSON, no markdown:\n"
            '{\n'
            '  "procurement_request": "<1-2 sentence professional description>",\n'
            '  "budget_limit": <number, monthly budget, default 5000 if not mentioned>,\n'
            '  "confidence": "<high|medium|low>"\n'
            '}'
        )

        try:
            content = _llm_call(self.llm, prompt)
            parsed  = json.loads(content)
            req     = parsed.get("procurement_request") or raw
            budget  = float(parsed.get("budget_limit") or state.get("budget_limit") or 5000)
            conf    = parsed.get("confidence", "medium")
            log    += f" | Request: '{req[:60]}' | Budget: ${budget:,.0f} | Confidence: {conf}"
        except Exception as e:
            # Fallback: raw text, keep existing budget
            req    = raw
            budget = float(state.get("budget_limit") or 5000)
            log   += f" | Parse failed ({e}), using raw input verbatim."

        return Command(
            update={"procurement_request": req, "budget_limit": budget, "logs": [log]},
            goto="research_node",
        )

    # ── Agent 1: Research ─────────────────────────────────────────────────────
    def research_node(self, state: ProcurementState) -> Command:
        request = (state.get("procurement_request") or "").strip()
        ts      = datetime.now().strftime("%H:%M:%S")

        if not request:
            log = f"[{ts}] ResearchNode: No request text — aborting."
            return Command(update={"logs": [log]}, goto=END)

        vendors = mock_vendor_search(request)
        log     = f"[{ts}] ResearchNode: Found {len(vendors)} vendor(s) for '{request[:50]}...'"

        return Command(
            update={"vendor_options": vendors, "logs": [log]},
            goto="analysis_node",
        )

    # ── Agent 2: Financial Analysis ───────────────────────────────────────────
    def analysis_node(self, state: ProcurementState) -> Command:
        vendors = state.get("vendor_options") or []
        budget  = float(state.get("budget_limit") or 5000)
        ts      = datetime.now().strftime("%H:%M:%S")
        log     = f"[{ts}] AnalysisNode: Evaluating {len(vendors)} vendor(s) vs budget ${budget:,.0f}"

        if not vendors:
            log += " | No vendors to evaluate."
            return Command(update={"analysis_approved": False, "logs": [log]}, goto=END)

        prompt = (
            f"You are a financial analyst. Budget limit: ${budget}/month.\n"
            f"Vendors:\n{json.dumps(vendors, indent=2)}\n\n"
            "Pick the BEST vendor whose price_per_month is AT OR BELOW the budget. "
            "If none fit, pick the cheapest one.\n"
            "Reply with ONLY valid JSON:\n"
            '{"recommended_vendor_id": "<vendor_id>", "rationale": "<one sentence>"}'
        )

        try:
            content   = _llm_call(self.llm, prompt)
            parsed    = json.loads(content)
            rec_id    = parsed.get("recommended_vendor_id", "")
            # Find the recommended vendor; fallback to cheapest
            selected  = min(vendors, key=lambda v: v["price_per_month"])
            for v in vendors:
                if v["vendor_id"] == rec_id:
                    selected = v
                    break
        except Exception:
            selected = min(vendors, key=lambda v: v["price_per_month"])

        result = validate_budget(selected["price_per_month"], budget)
        log   += f" | Selected: {selected['name']} (${selected['price_per_month']:,}/mo) | {result['status']}"

        if result["within_budget"]:
            return Command(
                update={"selected_vendor": selected, "analysis_approved": True, "logs": [log]},
                goto="approval_gate",
            )
        else:
            log += " | All vendors exceed budget — workflow rejected."
            return Command(
                update={"selected_vendor": selected, "analysis_approved": False, "logs": [log]},
                goto=END,
            )

    # ── Agent 3: Legal ────────────────────────────────────────────────────────
    def legal_node(self, state: ProcurementState) -> Command:
        vendor  = state.get("selected_vendor") or {}
        request = state.get("procurement_request") or ""
        ts      = datetime.now().strftime("%H:%M:%S")
        log     = f"[{ts}] LegalNode: Drafting contract for {vendor.get('name', 'Unknown')}"

        prompt = (
            "Draft a professional PURCHASE AGREEMENT in plain text (no markdown headers).\n\n"
            f"Vendor: {vendor.get('name', 'N/A')}\n"
            f"Service: {vendor.get('service_type', 'N/A')}\n"
            f"Monthly Cost: ${vendor.get('price_per_month', 0):,}\n"
            f"Capabilities: {', '.join(vendor.get('capabilities', []))}\n"
            f"Contract Terms: {vendor.get('contract_terms', 'Standard')}\n"
            f"Procurement Need: {request}\n\n"
            "Include: Parties, Services, Pricing & Payment, Term & Termination, "
            "Warranties, Confidentiality, Signatures."
        )

        contract = _llm_call(self.llm, prompt)
        log     += " | Contract generated successfully."

        return Command(
            update={"contract_draft": contract, "logs": [log]},
            goto=END,
        )


# =============================================================================
# 5. APPROVAL GATE & ROUTER
# =============================================================================

def approval_gate(state: ProcurementState) -> dict:
    vendor_name = (state.get("selected_vendor") or {}).get("name", "Unknown")
    ts  = datetime.now().strftime("%H:%M:%S")
    log = f"[{ts}] ApprovalGate: Paused — awaiting human decision for '{vendor_name}'"
    return {"logs": [log]}


def _approval_router(state: ProcurementState) -> str:
    approved = state.get("human_approved")
    return "legal_node" if approved else END


# =============================================================================
# 6. GRAPH
# =============================================================================

def build_procurement_graph(llm: ChatGroq) -> tuple:
    agents = ProcurementAgents(llm)
    graph  = StateGraph(ProcurementState)

    graph.add_node("parse_node",    agents.parse_node)
    graph.add_node("research_node", agents.research_node)
    graph.add_node("analysis_node", agents.analysis_node)
    graph.add_node("approval_gate", approval_gate)
    graph.add_node("legal_node",    agents.legal_node)

    graph.add_edge(START, "parse_node")
    # parse_node / research_node / analysis_node route via Command.goto
    graph.add_conditional_edges("approval_gate", _approval_router)
    graph.add_edge("legal_node", END)

    memory   = MemorySaver()
    compiled = graph.compile(checkpointer=memory, interrupt_before=["legal_node"])
    return compiled, memory, agents


# =============================================================================
# 7. EXECUTOR
# =============================================================================

class ProcurementWorkflowExecutor:
    def __init__(self, graph, checkpointer):
        self.graph        = graph
        self.checkpointer = checkpointer
        # UUID ensures no thread_id collisions across rapid calls
        self.thread_id    = f"proc_{uuid.uuid4().hex[:12]}"

    def _config(self) -> dict:
        return {"configurable": {"thread_id": self.thread_id}}

    def _safe_state(self) -> dict:
        snap = self.graph.get_state(self._config())
        if snap is None or snap.values is None:
            return _empty_state()
        return dict(snap.values)

    def start_workflow(self, procurement_request: str, budget_limit: float, raw_input: str = "") -> tuple:
        initial = {
            **_empty_state(),
            "raw_input":           raw_input.strip(),
            "procurement_request": procurement_request.strip(),
            "budget_limit":        float(budget_limit),
            "logs":                [f"Workflow started at {datetime.now().strftime('%H:%M:%S')}"],
        }
        events = []
        for event in self.graph.stream(initial, self._config(), stream_mode="updates"):
            events.append(event)
        return self._safe_state(), events

    def get_state(self) -> dict:
        return self._safe_state()

    def approve_vendor(self, approval: bool) -> tuple:
        self.graph.update_state(self._config(), {"human_approved": approval})
        events = []
        for event in self.graph.stream(None, self._config(), stream_mode="updates"):
            events.append(event)
        return self._safe_state(), events


# =============================================================================
# 8. LLM FACTORY
# =============================================================================

def create_llm() -> ChatGroq:
    return ChatGroq(
        model=GROQ_MODEL,
        api_key=GROQ_API_KEY,
        temperature=0.3,
        max_tokens=4096,
    )