# ── CELL 1: Imports & API Keys ──────────────────────────────────────────────
# Uses a local Ollama model — no API key needed.
# Make sure Ollama is installed and running, and the model is pulled:
#   ollama pull qwen2.5:7b-instruct

import os
import time
import uuid
from typing import TypedDict
from dotenv import load_dotenv

from ddgs import DDGS                                          # DuckDuckGo search
from langchain_huggingface import HuggingFaceEndpoint                   # Local Ollama LLM (replaces Gemini)
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import interrupt, Command

load_dotenv()


# ── Logging helper: BACKEND ONLY — this prints to the terminal/server console.
# It is intentionally NOT surfaced in the Streamlit UI; the UI only shows a
# generic "running" status. Watch your terminal to see step-by-step progress.
def log(msg: str):
    print(msg)


# ── CELL 2: State ────────────────────────────────────────────────────────────
class ResearchState(TypedDict):
    query:          str   # Original question from the user
    web_results:    str   # Formatted text of DuckDuckGo search results
    search_query:   str   # Rewritten / improved search query
    final_answer:   str   # LLM-synthesized answer
    needs_retry:    bool  # True if critic says answer is poor
    retry_count:    int   # Number of retry attempts so far
    human_approved: bool  # Whether the human approved the answer
    human_feedback: str   # Feedback from the human if they rejected


# ── CELL 3: LLM (Ollama, local) ──────────────────────────────────────────────
# qwen2.5:7b-instruct is a strong general-purpose pick for this pipeline:
# it follows formatting instructions well (Markdown, ## headings, [1]/[2]
# citations), is a reliable yes/no judge for the critic step, and is small
# enough (7B, ~4-5GB) to run comfortably on a single consumer GPU or even CPU.
# Override via env vars if you want a different model/host.
OLLAMA_MODEL    = os.environ.get("OLLAMA_MODEL", "qwen2.5:7b-instruct")
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")

llm = HuggingFaceEndpoint(
    repo_id="Qwen/Qwen2.5-7B-Instruct",
    temperature=0.3,
    huggingfacehub_api_token=os.environ["HF_TOKEN"]
)


# ── Retry helper: handles transient local errors (Ollama not up yet, model ──
# still loading into memory, brief timeout) with a short exponential backoff.
# Unlike a hosted API, Ollama has no request quota — retries here are just
# for "server not ready yet", not rate limiting.
def llm_invoke_with_retry(messages, max_retries: int = 3, base_wait: float = 3.0):
    """
    Calls llm.invoke(messages). If a connection/timeout error occurs (e.g.
    `ollama serve` isn't running yet, or the model is still loading),
    waits and retries up to max_retries times with exponential backoff.
    Raises the last exception if all retries are exhausted.
    """
    for attempt in range(max_retries + 1):
        try:
            if attempt == 0:
                log(f"    📡 Calling Ollama model '{OLLAMA_MODEL}'...")
            return llm.invoke(messages)
        except Exception as e:
            err_str = str(e)
            is_transient = any(s in err_str.lower() for s in
                                ["connection", "timeout", "timed out", "refused"])
            if is_transient and attempt < max_retries:
                wait_time = base_wait * (2 ** attempt)   # 3s, 6s, 12s
                log(f"    ⚠️  Ollama not responding yet. Waiting {wait_time:.0f}s "
                    f"before retry {attempt + 1}/{max_retries}... "
                    f"(is `ollama serve` running and is '{OLLAMA_MODEL}' pulled?)")
                time.sleep(wait_time)
                log(f"    📡 Retrying Ollama call ({attempt + 1}/{max_retries})...")
            else:
                log(f"    ❌ Ollama call failed permanently: {err_str[:200]}")
                raise


# ── CELL 4: Node 1 — Query Rewrite ───────────────────────────────────────────
REWRITE_PROMPT = """\
You rewrite user questions into a short, high-signal web search query.
Return ONLY the search query text. No explanation, no quotes.
"""

def query_rewrite_node(state: ResearchState) -> dict:
    log("🔄 STEP: Query Rewrite — turning your question into a search query")
    feedback = state.get("human_feedback", "").strip()

    if feedback:
        user_content = (
            f"Original question: {state['query']}\n"
            f"Previous answer was rejected. Human feedback: {feedback}\n"
            f"Rewrite the search query to address this feedback."
        )
        log(f"    Using human feedback to improve query: '{feedback}'")
    else:
        user_content = state["query"]

    messages = [
        SystemMessage(content=REWRITE_PROMPT),
        HumanMessage(content=user_content)
    ]
    response = llm_invoke_with_retry(messages)
    rewritten = response.content.strip()

    log(f"    ✅ Search query: {rewritten}")
    return {"search_query": rewritten}


# ── CELL 5: Node 2 — Web Search (DuckDuckGo) ─────────────────────────────────
def web_search_node(state: ResearchState) -> dict:
    log(f"🌐 STEP: Web Search — querying DuckDuckGo for '{state['search_query']}'")
    formatted_results = []

    try:
        with DDGS() as ddgs:
            raw = list(ddgs.text(state["search_query"], max_results=5))
        log(f"    ✅ Got {len(raw)} results")
    except Exception as e:
        log(f"    ❌ DuckDuckGo search error: {e}")
        raw = []

    for i, result in enumerate(raw, start=1):
        title   = result.get("title",  "No title")
        url     = result.get("href",   "")          # DuckDuckGo uses 'href' not 'url'
        content = result.get("body",   "").strip()  # DuckDuckGo uses 'body' not 'content'
        formatted_results.append(
            f"[{i}] {title}\n"
            f"    URL: {url}\n"
            f"    {content}"
        )

    web_results_text = "\n\n".join(formatted_results)
    return {"web_results": web_results_text}


# ── CELL 6: Node 3 — Synthesizer (Ollama) ────────────────────────────────────
SYSTEM_PROMPT = """\
You are an expert research assistant.

You will be given:
  1. A user research query.
  2. Five web search results (title, URL, snippet).

Your job:
  - Synthesize the web results with your own knowledge.
  - Write a clear, well-structured, comprehensive answer using Markdown.
  - Use ## headings to organize sections (e.g. ## Overview, ## Key Findings).
  - Cite sources inline using [1], [2], etc. where relevant.
  - Highlight any conflicting information across sources.
  - End with a "## References" section listing all URLs as:
      [1] https://...
      [2] https://...
  - Be factual. Do not hallucinate beyond the provided data.
"""

def synthesizer_node(state: ResearchState) -> dict:
    log("🧠 STEP: Synthesizer — asking the local Ollama model to write the sourced answer")
    user_prompt = (
        f"Research Query: {state['query']}\n\n"
        f"Web Search Results:\n\n{state['web_results']}"
    )

    messages = [
        SystemMessage(content=SYSTEM_PROMPT),
        HumanMessage(content=user_prompt)
    ]

    response = llm_invoke_with_retry(messages)
    log("    ✅ Answer drafted")
    return {"final_answer": response.content}


# ── CELL 7: Node 4 — Critic ──────────────────────────────────────────────────
CRITIC_PROMPT = """\
You judge whether a research answer is well supported by the given sources.
Reply with exactly one word: "yes" if well supported,
"no" if it is thin, vague, or barely uses the sources.
"""

def critic_node(state: ResearchState) -> dict:
    log("🔍 STEP: Critic — checking whether the answer is well supported")
    user_prompt = (
        f"Sources:\n{state['web_results']}\n\n"
        f"Answer:\n{state['final_answer']}"
    )

    messages = [
        SystemMessage(content=CRITIC_PROMPT),
        HumanMessage(content=user_prompt)
    ]

    verdict = llm_invoke_with_retry(messages).content.strip().lower()

    if verdict.startswith("no") and state["retry_count"] < 3:
        needs_retry = True
        log(f"    🔁 Verdict: '{verdict}' — retrying search (attempt {state['retry_count'] + 1}/3)")
    else:
        needs_retry = False
        log(f"    ✅ Verdict: '{verdict}' — moving to human review")

    return {
        "needs_retry": needs_retry,
        "retry_count": state["retry_count"] + 1
    }


# ── CELL 8: Node 5 — Human Approval ──────────────────────────────────────────
def human_approval_node(state: ResearchState) -> dict:
    log("✅ STEP: Human Approval — pausing for your review")

    # Pause graph — resumes when user submits feedback via Streamlit
    user_input = interrupt("Type 'ok' to approve, or give feedback to improve the answer: ")

    approved = user_input.strip().lower() == "ok"

    if approved:
        log("✅ Approved!")
        return {
            "human_approved": True,
            "human_feedback": ""
        }
    else:
        feedback = user_input.strip()
        log(f"🔄 Rejected. Feedback: '{feedback}'. Retrying with improved query...")
        return {
            "human_approved": False,
            "human_feedback": feedback,
            "retry_count":    0
        }


# ── CELL 9: Build Graph ───────────────────────────────────────────────────────
#
#   START → query_rewrite → web_search → synthesizer → critic
#               ↑                                         |
#               └──── needs_retry=True ───────────────────┘
#                              (else)
#                                ↓
#                        human_approval   ← graph PAUSES here
#                        /           \
#                  approved         rejected (with feedback)
#                     ↓                  ↓
#                    END          query_rewrite  (smarter retry)

def route_after_critic(state: ResearchState):
    return "web_search" if state["needs_retry"] else "human_approval"

def route_after_human(state: ResearchState):
    return "end" if state["human_approved"] else "query_rewrite"


workflow = StateGraph(ResearchState)

workflow.add_node("query_rewrite",  query_rewrite_node)
workflow.add_node("web_search",     web_search_node)
workflow.add_node("synthesizer",    synthesizer_node)
workflow.add_node("critic",         critic_node)
workflow.add_node("human_approval", human_approval_node)

workflow.add_edge(START,           "query_rewrite")
workflow.add_edge("query_rewrite", "web_search")
workflow.add_edge("web_search",    "synthesizer")
workflow.add_edge("synthesizer",   "critic")

workflow.add_conditional_edges(
    "critic",
    route_after_critic,
    {"web_search": "web_search", "human_approval": "human_approval"}
)

workflow.add_conditional_edges(
    "human_approval",
    route_after_human,
    {"end": END, "query_rewrite": "query_rewrite"}
)

memory = MemorySaver()
agent  = workflow.compile(checkpointer=memory)


# ── Helper: create a fresh thread config ─────────────────────────────────────
def new_thread() -> dict:
    return {"configurable": {"thread_id": str(uuid.uuid4())}}