from __future__ import annotations

import json
import re
from typing import Annotated, Literal
from typing_extensions import TypedDict

from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver

# ── State schema ──────────────────────────────────────────────
# add_messages is a reducer: it appends new messages instead of replacing.
# All other fields use last-write-wins (default for TypedDict in LangGraph).
class PRReviewState(TypedDict):
    messages:          Annotated[list, add_messages]
    pr_diff:           str
    test_results:      str
    security_findings: list[dict]   # [{severity, category, line, detail}]
    risk_score:        float         # 0.0 (clean) → 1.0 (critical block)
    recommendation:    str           # "approve" | "request_changes" | "escalate"

# ── Tools (functions the LLM can call via tool_use) ───────────
@tool
def run_semgrep(diff: str) -> list[dict]:
    """Run static SAST analysis on a unified diff. Returns SARIF-lite findings."""
    # In production: subprocess.run(["semgrep", "--config=auto", "--json", ...]).
    # Here we return a typed stub so the graph structure is realistic.
    return [{"severity": "WARNING", "category": "injection",
             "line": 42, "detail": "Potential SQL injection via f-string"}]

@tool
def get_ci_failures(pr_number: int) -> list[str]:
    """Fetch failing test names from the last CI run for a PR."""
    # In production: GitHub API → GET /repos/{owner}/{repo}/commits/{sha}/check-runs
    return ["tests/test_auth.py::test_token_expiry"]

tools = [run_semgrep, get_ci_failures]

# ── LLM bound to tools ────────────────────────────────────────
llm = ChatAnthropic(
    model="claude-sonnet-4-6",
    max_tokens=1024,
).bind_tools(tools)

SYSTEM_PROMPT = """You are a senior security-aware code reviewer embedded in a CI/CD pipeline.
Your only actions are calling the provided tools and returning structured JSON.
Never produce narrative text. Always respond with a JSON object:
{"security_findings": [...], "risk_score": 0.0-1.0, "rationale": "..."}"""

# ── Graph nodes ───────────────────────────────────────────────
def analyse_diff(state: PRReviewState) -> dict:
    # Invoke LLM with tool access; it may call run_semgrep or get_ci_failures.
    response = llm.invoke([
        SystemMessage(content=SYSTEM_PROMPT),
        HumanMessage(content=f"Review this PR diff:\n\n{state['pr_diff']}"),
    ])
    return {"messages": [response]}

def parse_llm_output(state: PRReviewState) -> dict:
    # Extract structured fields from the LLM's last message.
    last_msg = state["messages"][-1]
    try:
        # Strip markdown fences the model may still emit despite the system prompt.
        raw = re.sub(r"```(?:json)?\n?|```", "", last_msg.content).strip()
        parsed = json.loads(raw)
        return {
            "security_findings": parsed.get("security_findings", []),
            "risk_score":        float(parsed.get("risk_score", 0.5)),
        }
    except (json.JSONDecodeError, KeyError, TypeError):
        # Parsing failure → conservative default: escalate to human
        return {"security_findings": [], "risk_score": 1.0}

def apply_decision(state: PRReviewState, rec: str) -> dict:
    return {"recommendation": rec}

# ── Conditional router ────────────────────────────────────────
def route_by_risk(state: PRReviewState) -> Literal["approve", "request_changes", "escalate"]:
    score = state["risk_score"]
    if score < 0.3:  return "approve"
    if score < 0.75: return "request_changes"
    return "escalate"

# ── Build and compile the graph ───────────────────────────────
builder = StateGraph(PRReviewState)

builder.add_node("analyse",        analyse_diff)
builder.add_node("parse",          parse_llm_output)
builder.add_node("approve",        lambda s: apply_decision(s, "approve"))
builder.add_node("request_changes", lambda s: apply_decision(s, "request_changes"))
builder.add_node("escalate",       lambda s: apply_decision(s, "escalate"))

builder.set_entry_point("analyse")
builder.add_edge("analyse", "parse")
builder.add_conditional_edges("parse", route_by_risk)
for terminal in ("approve", "request_changes", "escalate"):
    builder.add_edge(terminal, END)

# MemorySaver checkpoints state after each node - enables time-travel debugging.
agent = builder.compile(checkpointer=MemorySaver())

# ── Usage ─────────────────────────────────────────────────────
if __name__ == "__main__":
    pr_diff   = "+ cursor.execute(f'SELECT * FROM users WHERE id={user_id}')"
    ci_output = "FAILED tests/test_auth.py::test_token_expiry"

    result = agent.invoke(
        {"pr_diff": pr_diff, "test_results": ci_output, "messages": []},
        config={"configurable": {"thread_id": "pr-1042"}},   # thread_id = checkpoint key
    )
    print(result["recommendation"])   # → "escalate" (SQL injection, failing test)
    print(result["risk_score"])        # → 1.0

import time, uuid
import structlog
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from anthropic import AsyncAnthropic

log = structlog.get_logger()
app = FastAPI(title="AI Gateway")
client = AsyncAnthropic()

# Per-agent token budget (tokens/minute). Enforced at the gateway, not in the agent.
AGENT_BUDGETS: dict[str, int] = {
    "pr-review-agent":      50_000,
    "incident-triage-agent": 30_000,
    "log-analysis-agent":    20_000,
}
_usage: dict[str, tuple[int, float]] = {}  # {agent_id: (tokens_used, window_start)}

class GatewayRequest(BaseModel):
    agent_id:  str
    model:     str = "claude-sonnet-4-6"
    messages:  list[dict]
    max_tokens: int = 1024
    system:    str | None = None

@app.post("/v1/messages")
async def proxy_to_anthropic(req: GatewayRequest) -> JSONResponse:
    correlation_id = str(uuid.uuid4())
    budget = AGENT_BUDGETS.get(req.agent_id)
    if budget is None:
        raise HTTPException(403, f"Unknown agent_id: {req.agent_id}")

    # Sliding-window rate limiter (60-second window per agent)
    used, window_start = _usage.get(req.agent_id, (0, time.monotonic()))
    if time.monotonic() - window_start > 60:
        used, window_start = 0, time.monotonic()  # reset window
    if used >= budget:
        log.warning("agent_budget_exceeded", agent_id=req.agent_id, used=used, budget=budget)
        raise HTTPException(429, "Agent token budget exceeded for this window")

    t0 = time.perf_counter()
    response = await client.messages.create(
        model=req.model, messages=req.messages,
        max_tokens=req.max_tokens,
        **({"system": req.system} if req.system else {}),
    )
    latency_ms = (time.perf_counter() - t0) * 1000

    tokens_used = response.usage.input_tokens + response.usage.output_tokens
    _usage[req.agent_id] = (used + tokens_used, window_start)

    # Structured audit log - every agent decision is traceable by correlation_id.
    log.info("agent_llm_call",
        correlation_id=correlation_id,
        agent_id=req.agent_id,
        model=req.model,
        input_tokens=response.usage.input_tokens,
        output_tokens=response.usage.output_tokens,
        latency_ms=round(latency_ms, 1),
        stop_reason=response.stop_reason,
    )
    return JSONResponse(response.model_dump())

from pydantic import BaseModel, Field, field_validator
from langchain_anthropic import ChatAnthropic
from langchain_core.messages import HumanMessage, SystemMessage

# Structured output forces the judge to commit to a verdict - no hedging.
class JudgeVerdict(BaseModel):
    valid:           bool
    confidence:      float = Field(ge=0.0, le=1.0)
    override_reason: str | None = None
    final_risk:      float = Field(ge=0.0, le=1.0)

    @field_validator("confidence")
    def confidence_not_zero(cls, v: float) -> float:
        if v == 0.0:
            raise ValueError("Judge must commit to a non-zero confidence")
        return v

JUDGE_SYSTEM = """You are a senior security reviewer validating an AI agent's PR analysis.
You receive: (1) the original diff, (2) the agent's findings, (3) the agent's risk score.
Return a JSON object matching this schema exactly:
{"valid": bool, "confidence": 0.0-1.0, "override_reason": str|null, "final_risk": 0.0-1.0}
Rules:
- valid=true if the agent's findings are supported by the diff content.
- Set override_reason if you are adjusting the risk score; explain why.
- If you cannot evaluate the diff, set valid=false and final_risk=1.0 (fail safe)."""

# Use a lower temperature for the judge - we want deterministic evaluation, not creativity.
judge_llm = ChatAnthropic(model="claude-sonnet-4-6", temperature=0.1)

def llm_judge(state: PRReviewState) -> dict:
    """LangGraph node: validate the primary agent's output before routing."""
    verdict_msg = judge_llm.invoke([
        SystemMessage(content=JUDGE_SYSTEM),
        HumanMessage(content=json.dumps({
            "diff":             state["pr_diff"][:3000],  # cap context to stay in budget
            "agent_findings":   state["security_findings"],
            "agent_risk_score": state["risk_score"],
        })),
    ])
    try:
        raw = re.sub(r"```(?:json)?\n?|```", "", verdict_msg.content).strip()
        verdict = JudgeVerdict.model_validate_json(raw)
        return {"risk_score": verdict.final_risk}  # overwrite with judge's adjusted score
    except Exception:
        return {"risk_score": 1.0}  # judge failure → conservative escalation

# Insert the judge node between parse and the routing decision:
# analyse → parse → judge → route_by_risk → (approve|request_changes|escalate)
builder.add_node("judge", llm_judge)
builder.add_edge("parse", "judge")
builder.add_conditional_edges("judge", route_by_risk)