from typing import TypedDict, Annotated, Literal
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.types import interrupt
import operator

# 1. Define state schema — typed, explicit
class PRReviewState(TypedDict):
    messages:      Annotated[list, operator.add]
    pr_url:        str
    code_issues:   list[str]
    security_flags: list[str]
    test_coverage: float
    approved:      bool
    reviewer_notes: str

# 2. Define nodes — each is a pure function (state in, updates out)
def analyse_code(state: PRReviewState) -> dict:
    """Node: LLM analyses the code changes"""
    response = llm_with_tools.invoke([
        SystemMessage(content=CODE_REVIEW_PROMPT),
        HumanMessage(content=f"Review PR: {state['pr_url']}")
    ])
    return {"messages": [response]}

def security_scan(state: PRReviewState) -> dict:
    """Node: Dedicated security analysis"""
    code = fetch_pr_diff(state["pr_url"])
    flags = run_security_tools(code)  # Bandit, Semgrep, etc.
    return {"security_flags": flags}

def human_approval(state: PRReviewState) -> dict:
    """Node: Pause and wait for human engineer approval"""
    decision = interrupt({
        "action": "approve_pr_comment",
        "pr_url": state["pr_url"],
        "code_issues": state["code_issues"],
        "security_flags": state["security_flags"],
        "message": "Review agent findings and approve posting to GitHub?"
    })
    return {
        "approved": decision["approved"],
        "reviewer_notes": decision.get("notes", "")
    }

def post_review(state: PRReviewState) -> dict:
    """Node: Post approved review to GitHub"""
    post_github_comment(state["pr_url"], build_review_comment(state))
    return {"messages": [SystemMessage(content="Review posted to GitHub.")]}

# 3. Define routing logic
def route_after_analysis(state: PRReviewState) -> Literal["tools", "security_scan"]:
    last = state["messages"][-1]
    return "tools" if hasattr(last, "tool_calls") and last.tool_calls else "security_scan"

def route_after_approval(state: PRReviewState) -> Literal["post_review", END]:
    return "post_review" if state["approved"] else END

# 4. Build and compile the graph
graph = StateGraph(PRReviewState)
graph.add_node("analyse_code", analyse_code)
graph.add_node("tools", ToolNode(tools))
graph.add_node("security_scan", security_scan)
graph.add_node("human_approval", human_approval)
graph.add_node("post_review", post_review)

graph.set_entry_point("analyse_code")
graph.add_conditional_edges("analyse_code", route_after_analysis)
graph.add_edge("tools", "analyse_code")
graph.add_edge("security_scan", "human_approval")
graph.add_conditional_edges("human_approval", route_after_approval)
graph.add_edge("post_review", END)

# 5. Compile with production checkpointer
checkpointer = PostgresSaver.from_conn_string("postgresql://...")
app = graph.compile(checkpointer=checkpointer, interrupt_before=["human_approval"])

# 6. Run with persistent thread
config = {"configurable": {"thread_id": f"pr-review-{pr_number}"}}
result = app.invoke({"pr_url": pr_url, "messages": []}, config=config)

# Stream every event from the graph in real time
async for event in app.astream_events(inputs, config=config, version="v2"):
    kind = event["event"]

    if kind == "on_chat_model_stream":
        # Token-by-token streaming of LLM output
        chunk = event["data"]["chunk"]
        print(chunk.content, end="", flush=True)

    elif kind == "on_tool_start":
        # Tool about to be called
        print(f"\n→ Calling: {event['name']}({event['data']['input']})")

    elif kind == "on_tool_end":
        # Tool completed
        print(f"← Result: {str(event['data']['output'])[:100]}...")

    elif kind == "on_chain_end" and event["name"] == "LangGraph":
        # Full graph execution complete
        final_state = event["data"]["output"]

from crewai import Agent, Task, Crew, Process
from crewai_tools import SerperDevTool, ScrapeWebsiteTool
from langchain_groq import ChatGroq

# Free LLM via Groq
llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0.1)
search = SerperDevTool()
scraper = ScrapeWebsiteTool()

# ── DEFINE AGENTS (specialists with distinct roles) ───────────

researcher = Agent(
    role="Senior Market Research Analyst",
    goal="Find accurate, comprehensive information about {company}. Never invent facts.",
    backstory="""You are a veteran market analyst with 15 years at Goldman Sachs.
    You are known for rigorous fact-checking and citing every claim.
    You dig deep — finding information others miss.""",
    tools=[search, scraper],
    llm=llm,
    verbose=True,
    max_iter=10
)

financial_analyst = Agent(
    role="Financial Intelligence Analyst",
    goal="Analyse funding, revenue, valuation, and financial health of {company}",
    backstory="""You are a CFA charterholder who has evaluated 500+ companies.
    You spot financial red flags others miss and always triangulate data
    from multiple sources before forming conclusions.""",
    tools=[search],
    llm=llm,
    verbose=True
)

risk_analyst = Agent(
    role="Risk & Competitive Intelligence Analyst",
    goal="Identify risks, competitive threats, and regulatory issues for {company}",
    backstory="""You have worked in risk management at BlackRock.
    You systematically assess market risks, competitive dynamics,
    regulatory exposure, and execution risks for any company.""",
    tools=[search],
    llm=llm,
    verbose=True
)

report_writer = Agent(
    role="Investment Memo Writer",
    goal="Synthesise research into a clear, structured investment memo",
    backstory="""You are a former managing director who has written
    thousands of investment memos. You know how to make complex
    information clear, structured, and actionable for decision-makers.""",
    llm=llm,
    verbose=True
)

# ── DEFINE TASKS (what each agent produces) ───────────────────

research_task = Task(
    description="""Research {company} thoroughly:
    1. What exactly does the company do? Business model in detail.
    2. Who are the founders and key executives?
    3. What products/services do they offer? Recent launches?
    4. Any recent news (last 6 months)?
    5. Their main customers and market segment?
    Cite every source with URL.""",
    expected_output="Detailed research notes with source URLs for every fact",
    agent=researcher
)

financial_task = Task(
    description="""Analyse {company} financial profile:
    1. All funding rounds — dates, amounts, investors, valuation
    2. Revenue if available (public company or leaked data)
    3. Burn rate estimates and runway
    4. Valuation multiples vs comparable companies
    5. Last 3 financial milestones announced
    Only include verified figures — mark anything uncertain as [ESTIMATED].""",
    expected_output="Financial analysis with all funding data, valuations, and comparables",
    agent=financial_analyst,
    context=[research_task]  # Gets researcher output as input
)

risk_task = Task(
    description="""Identify key risks for {company}:
    1. Top 3 competitive threats and who poses them
    2. Regulatory risks in their market
    3. Execution risks (team gaps, scaling challenges)
    4. Market risks (timing, macro environment)
    5. Any reported controversies or legal issues
    Rate each risk: High / Medium / Low.""",
    expected_output="Risk assessment with H/M/L ratings and supporting evidence",
    agent=risk_analyst,
    context=[research_task, financial_task]
)

memo_task = Task(
    description="""Write a professional investment memo for {company}.
    Use all research, financial, and risk data provided.
    Structure:
    # {company} — Investment Memo
    ## Executive Summary (3 sentences max)
    ## Company Overview
    ## Business Model & Revenue
    ## Financial Profile
    ## Competitive Landscape
    ## Risk Assessment
    ## Investment Thesis
    ## Recommendation: [INVEST / PASS / MONITOR]
    ## Sources""",
    expected_output="Complete investment memo in markdown format",
    agent=report_writer,
    context=[research_task, financial_task, risk_task],
    output_file="investment_memo_{company}.md"
)

# ── ASSEMBLE AND RUN CREW ─────────────────────────────────────

crew = Crew(
    agents=[researcher, financial_analyst, risk_analyst, report_writer],
    tasks=[research_task, financial_task, risk_task, memo_task],
    process=Process.sequential,  # Or Process.hierarchical for manager agent
    verbose=True,
    memory=True,               # Enable crew-level memory (Qdrant backed)
    embedder={"provider": "openai", "config": {"model": "text-embedding-3-small"}}
)

result = crew.kickoff(inputs={"company": "Adyen"})
print(result.raw)

# Manager agent routes tasks to specialists
manager = Agent(
    role="Research Director",
    goal="Coordinate the research team to produce the best possible output",
    backstory="You manage a team of research specialists and know exactly which tasks to assign to whom.",
    llm=llm,
    allow_delegation=True  # This agent can delegate to others
)

crew = Crew(
    agents=[researcher, financial_analyst, risk_analyst, report_writer],
    tasks=[research_task, financial_task, risk_task, memo_task],
    process=Process.hierarchical,
    manager_agent=manager,      # Manager routes and coordinates
    verbose=True
)
# Manager sees all tasks, assigns them, reviews outputs, requests revisions

from agents import Agent, Runner, handoff, tool
import asyncio

# Specialist agents
billing_agent = Agent(
    name="Billing Specialist",
    instructions="""You handle all billing and payment questions.
    You have access to customer billing records and can process refunds up to $500.
    For refunds over $500, escalate to human support.""",
    tools=[get_invoice, process_refund, check_subscription]
)

technical_agent = Agent(
    name="Technical Support Specialist",
    instructions="""You handle technical issues, bugs, and integration questions.
    You have access to the customer's account config and recent error logs.
    If the issue requires a code fix, create a Jira ticket.""",
    tools=[get_error_logs, get_account_config, create_jira_ticket]
)

# Triage agent — routes to specialists
triage_agent = Agent(
    name="Customer Support Triage",
    instructions="""You are the first point of contact for customer support.
    Greet the customer warmly, understand their issue fully, then route:
    - Billing questions → Billing Specialist
    - Technical issues → Technical Support Specialist
    - General questions → answer directly
    Never guess — if unsure, ask a clarifying question first.""",
    handoffs=[
        handoff(billing_agent,   description="Route billing/payment issues here"),
        handoff(technical_agent, description="Route technical/integration issues here")
    ]
)

# Run the triage agent
async def handle_customer(message: str) -> str:
    result = await Runner.run(
        triage_agent,
        input=message,
        max_turns=20
    )
    return result.final_output

# Trace shows every step: which agent ran, what tools were called, when handoffs occurred

from pydantic_ai import Agent
from pydantic import BaseModel, Field
from typing import Literal

# Define structured output — agent MUST return this shape
class CompanyAnalysis(BaseModel):
    company_name: str = Field(description="Exact legal name")
    industry: str = Field(description="Primary industry vertical")
    founded_year: int = Field(description="Year founded", ge=1800, le=2026)
    total_funding_usd: float = Field(description="Total funding raised in USD")
    valuation_usd: float | None = Field(description="Last known valuation, None if unknown")
    recommendation: Literal["invest", "pass", "monitor"]
    confidence: float = Field(description="Confidence score 0-1", ge=0, le=1)
    key_risks: list[str] = Field(description="Top 3-5 risks", max_items=5)
    sources: list[str] = Field(description="URLs of sources used")

# Agent is forced to produce valid CompanyAnalysis — hallucinated fields are impossible
agent = Agent(
    "groq:llama-3.3-70b-versatile",
    result_type=CompanyAnalysis,
    system_prompt="""You are an investment analyst. Research companies thoroughly
    and produce structured analysis. Only include verified facts.
    For unknown values, use None rather than guessing.""",
    tools=[web_search, fetch_url]
)

result = await agent.run("Analyse Adyen as a potential investment")
analysis: CompanyAnalysis = result.data  # Fully typed, validated
print(f"Recommendation: {analysis.recommendation} ({analysis.confidence:.0%} confidence)")

async def debate_answer(question: str) -> str:
    # Round 1: independent answers
    agent_a_answer, agent_b_answer = await asyncio.gather(
        agent_a.run(question),
        agent_b.run(question)
    )
    # Round 2: each critiques the other
    critique_a, critique_b = await asyncio.gather(
        agent_a.run(f"Critique this answer: {agent_b_answer}. What's missing or wrong?"),
        agent_b.run(f"Critique this answer: {agent_a_answer}. What's missing or wrong?")
    )
    # Round 3: judge synthesises
    return await judge_agent.run(
        f"Question: {question}\n"
        f"Answer A: {agent_a_answer}\n"
        f"Answer B: {agent_b_answer}\n"
        f"Critique of A: {critique_b}\n"
        f"Critique of B: {critique_a}\n"
        f"Synthesise the best answer incorporating valid critiques:"
    )