import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Annotated, Any, TypedDict

import faiss
import numpy as np
import streamlit as st
from duckduckgo_search import DDGS
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_groq import ChatGroq
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
from sentence_transformers import SentenceTransformer

# --- Settings ---
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
GROQ_MODEL = os.getenv("GROQ_MODEL", "openai/gpt-oss-120b")
EMBEDDING_MODEL = "all-MiniLM-L6-v2"
MAX_RESEARCH_ITERATIONS = 2

MEMORY_DIR = Path(__file__).parent / "data" / "memory"
MEMORY_DIR.mkdir(parents=True, exist_ok=True)
FAISS_INDEX_PATH = MEMORY_DIR / "research_memory.index"
FAISS_META_PATH = MEMORY_DIR / "research_memory_meta.json"

AGENT_COLORS = {
    "Researcher": "#4A90D9",
    "Critic": "#E67E22",
    "Synthesizer": "#9B59B6",
    "Writer": "#27AE60",
}


# --- Memory ---
class ResearchMemory:
    def __init__(self) -> None:
        self._encoder = SentenceTransformer(EMBEDDING_MODEL)
        self._dim = self._encoder.get_sentence_embedding_dimension()
        self._index: faiss.IndexFlatL2 | None = None
        self._metadata: list[dict[str, Any]] = []
        self._load()

    def _load(self) -> None:
        if FAISS_INDEX_PATH.exists() and FAISS_META_PATH.exists():
            self._index = faiss.read_index(str(FAISS_INDEX_PATH))
            self._metadata = json.loads(FAISS_META_PATH.read_text(encoding="utf-8"))
        else:
            self._index = faiss.IndexFlatL2(self._dim)
            self._metadata = []

    def _save(self) -> None:
        if self._index is not None:
            faiss.write_index(self._index, str(FAISS_INDEX_PATH))
        FAISS_META_PATH.write_text(
            json.dumps(self._metadata, indent=2, ensure_ascii=False),
            encoding="utf-8",
        )

    def _embed(self, text: str) -> np.ndarray:
        vec = self._encoder.encode([text], normalize_embeddings=True)
        return np.array(vec, dtype=np.float32)

    def store(self, topic: str, report: str) -> None:
        entry = {
            "topic": topic,
            "report": report,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        vector = self._embed(f"{topic}\n{report[:2000]}")
        self._index.add(vector)
        self._metadata.append(entry)
        self._save()

    def recall(self, query: str, k: int = 3) -> list[dict[str, Any]]:
        if not self._metadata or self._index is None or self._index.ntotal == 0:
            return []
        vector = self._embed(query)
        distances, indices = self._index.search(vector, min(k, self._index.ntotal))
        results = []
        for dist, idx in zip(distances[0], indices[0]):
            if idx < 0 or idx >= len(self._metadata):
                continue
            entry = dict(self._metadata[idx])
            entry["score"] = float(dist)
            results.append(entry)
        return results

    def list_sessions(self) -> list[dict[str, Any]]:
        return [
            {
                "topic": m["topic"],
                "timestamp": m["timestamp"],
                "preview": m["report"][:200] + "..." if len(m["report"]) > 200 else m["report"],
            }
            for m in reversed(self._metadata)
        ]

    def clear(self) -> None:
        self._index = faiss.IndexFlatL2(self._dim)
        self._metadata = []
        self._save()


# --- Search ---
def web_search(query: str, max_results: int = 5) -> list[dict[str, str]]:
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=max_results))
        return [
            {"title": r.get("title", ""), "snippet": r.get("body", ""), "url": r.get("href", "")}
            for r in results
        ]
    except Exception:
        return []


def format_search_results(results: list[dict[str, str]]) -> str:
    if not results:
        return "No web results found. Use your training knowledge and note assumptions."
    lines = []
    for i, r in enumerate(results, 1):
        lines.append(f"[{i}] {r['title']}\n    URL: {r['url']}\n    {r['snippet']}")
    return "\n\n".join(lines)


# --- LangGraph state & agents ---
class ResearchState(TypedDict):
    topic: str
    memory_context: str
    research_notes: str
    sources: list[str]
    critique: str
    critique_approved: bool
    synthesis: str
    report: str
    iteration: int
    messages: Annotated[list, add_messages]


def _llm() -> ChatGroq:
    return ChatGroq(groq_api_key=GROQ_API_KEY, model_name=GROQ_MODEL, temperature=0.3)


def _invoke(system: str, user: str) -> str:
    response = _llm().invoke([SystemMessage(content=system), HumanMessage(content=user)])
    return response.content.strip()


def researcher_node(state: ResearchState) -> dict:
    topic = state["topic"]
    search_results = web_search(f"{topic} research overview latest developments", max_results=6)
    sources = [r["url"] for r in search_results if r.get("url")]

    user_parts = [f"Topic: {topic}"]
    if state.get("memory_context"):
        user_parts.append(f"Relevant past research from memory:\n{state['memory_context']}")
    if state.get("research_notes") and state.get("critique"):
        user_parts.append(
            f"Previous research (needs improvement):\n{state['research_notes']}\n\n"
            f"Critic feedback to address:\n{state['critique']}"
        )
    user_parts.append(f"Web search results:\n{format_search_results(search_results)}")

    notes = _invoke(
        "You are a Research Agent. Gather comprehensive, factual research on the given topic. "
        "Structure your notes with: Key Concepts, Current State, Important Developments, "
        "Open Questions, and Data Points. Be thorough and cite sources by URL when available.",
        "\n\n".join(user_parts),
    )
    return {
        "research_notes": notes,
        "sources": sources,
        "messages": [AIMessage(content=f"[Researcher] Completed research on: {topic}")],
    }


def critic_node(state: ResearchState) -> dict:
    iteration = state.get("iteration", 0)
    critique = _invoke(
        "You are a Critic Agent. Rigorously evaluate research quality. "
        "Check for: factual gaps, weak reasoning, missing perspectives, outdated info, "
        "and lack of evidence. End your response with exactly one line: "
        "VERDICT: APPROVED or VERDICT: REVISE",
        f"Topic: {state['topic']}\n\nResearch notes to evaluate:\n{state['research_notes']}",
    )
    approved = "VERDICT: APPROVED" in critique.upper()
    force_approve = iteration >= MAX_RESEARCH_ITERATIONS - 1
    return {
        "critique": critique,
        "critique_approved": approved or force_approve,
        "iteration": iteration + 1,
        "messages": [
            AIMessage(
                content=f"[Critic] {'Approved' if (approved or force_approve) else 'Requested revision'}"
            )
        ],
    }


def synthesizer_node(state: ResearchState) -> dict:
    user_parts = [
        f"Topic: {state['topic']}",
        f"Research notes:\n{state['research_notes']}",
        f"Critic evaluation:\n{state['critique']}",
    ]
    if state.get("memory_context"):
        user_parts.append(f"Historical context from memory:\n{state['memory_context']}")

    synthesis = _invoke(
        "You are a Synthesizer Agent. Merge research findings and critique into a unified "
        "analytical framework. Identify themes, connections, contradictions, and insights. "
        "Produce a structured synthesis ready for report writing.",
        "\n\n".join(user_parts),
    )
    return {
        "synthesis": synthesis,
        "messages": [AIMessage(content="[Synthesizer] Unified analytical framework ready")],
    }


def writer_node(state: ResearchState) -> dict:
    sources = state.get("sources", [])
    source_block = "\n".join(f"- {url}" for url in sources) if sources else "No external URLs captured."
    report = _invoke(
        "You are a Writer Agent. Produce a polished, publication-quality deep-research report. "
        "Use markdown formatting with: Executive Summary, Introduction, Key Findings, "
        "Deep Analysis, Implications, Limitations, and Conclusion. "
        "Write in clear, authoritative prose. Include a Sources section at the end.",
        f"Topic: {state['topic']}\n\nSynthesis to expand into full report:\n{state['synthesis']}\n\n"
        f"Available sources:\n{source_block}",
    )
    return {
        "report": report,
        "messages": [AIMessage(content="[Writer] Deep-research report completed")],
    }


def _route_after_critic(state: ResearchState) -> str:
    return "synthesizer" if state.get("critique_approved", False) else "researcher"


def build_research_graph():
    graph = StateGraph(ResearchState)
    graph.add_node("researcher", researcher_node)
    graph.add_node("critic", critic_node)
    graph.add_node("synthesizer", synthesizer_node)
    graph.add_node("writer", writer_node)
    graph.set_entry_point("researcher")
    graph.add_edge("researcher", "critic")
    graph.add_conditional_edges(
        "critic",
        _route_after_critic,
        {"researcher": "researcher", "synthesizer": "synthesizer"},
    )
    graph.add_edge("synthesizer", "writer")
    graph.add_edge("writer", END)
    return graph.compile()


# --- Streamlit UI ---
st.set_page_config(page_title="Deep Research Agents", page_icon="🔬", layout="wide")


def init_session():
    if "memory" not in st.session_state:
        st.session_state.memory = ResearchMemory()
    if "graph" not in st.session_state:
        st.session_state.graph = build_research_graph()
    if "running" not in st.session_state:
        st.session_state.running = False


def render_agent_status(agent: str, status: str):
    color = AGENT_COLORS.get(agent, "#888")
    st.markdown(
        f'<div style="padding:8px 12px;margin:4px 0;border-left:4px solid {color};'
        f'background:#1a1a2e;border-radius:4px;">'
        f'<strong style="color:{color}">{agent}</strong> — {status}</div>',
        unsafe_allow_html=True,
    )


def run_research(topic: str, use_memory: bool):
    memory: ResearchMemory = st.session_state.memory
    graph = st.session_state.graph

    memory_context = ""
    if use_memory:
        recalled = memory.recall(topic, k=3)
        if recalled:
            memory_context = "\n\n---\n\n".join(
                f"Past topic: {r['topic']} (relevance score: {r['score']:.3f})\n{r['report'][:1500]}"
                for r in recalled
            )

    progress = st.progress(0, text="Initializing agents...")
    status_box = st.empty()
    stages = [
        ("Researcher", "Gathering information..."),
        ("Critic", "Evaluating research quality..."),
        ("Synthesizer", "Building analytical framework..."),
        ("Writer", "Composing final report..."),
    ]

    initial_state = {
        "topic": topic,
        "memory_context": memory_context,
        "research_notes": "",
        "sources": [],
        "critique": "",
        "critique_approved": False,
        "synthesis": "",
        "report": "",
        "iteration": 0,
        "messages": [],
    }

    final_state = None
    stage_idx = 0

    for event in graph.stream(initial_state, stream_mode="updates"):
        for node_name, node_output in event.items():
            stage_idx = {"researcher": 0, "critic": 1, "synthesizer": 2, "writer": 3}.get(node_name, stage_idx)
            progress.progress(min((stage_idx + 1) / len(stages), 1.0), text=stages[stage_idx][1])
            with status_box.container():
                for i, (agent, desc) in enumerate(stages):
                    if i <= stage_idx:
                        render_agent_status(agent, "✓ " + desc if i < stage_idx else desc)
            if final_state is None:
                final_state = {**initial_state, **node_output}
            else:
                final_state.update(node_output)

    progress.progress(1.0, text="Research complete!")
    return final_state


def main():
    init_session()
    memory: ResearchMemory = st.session_state.memory

    st.title("🔬 Deep Research Agents")
    st.caption(
        "Multi-agent research pipeline — **Researcher** → **Critic** → **Synthesizer** → **Writer** "
        "powered by Groq + LangGraph + FAISS"
    )

    with st.sidebar:
        st.header("Settings")
        use_memory = st.toggle("Use cross-session memory", value=True)
        st.divider()
        st.subheader("Past Sessions")
        sessions = memory.list_sessions()
        if sessions:
            for s in sessions[:10]:
                st.markdown(f"**{s['topic']}**")
                st.caption(s["timestamp"][:19])
                st.text(s["preview"][:120])
                st.divider()
        else:
            st.info("No past sessions yet.")
        if st.button("Clear memory", type="secondary"):
            memory.clear()
            st.rerun()

    if not GROQ_API_KEY:
        st.error(
            "Set your `GROQ_API_KEY` environment variable. "
            "Get a free key at [console.groq.com](https://console.groq.com)."
        )
        st.stop()

    col1, col2 = st.columns([3, 1])
    with col1:
        topic = st.text_input(
            "Research topic",
            placeholder="e.g. The impact of quantum computing on cryptography",
        )
    with col2:
        st.write("")
        st.write("")
        submit = st.button("Start Research", type="primary", disabled=st.session_state.running)

    if submit and topic.strip():
        st.session_state.running = True
        with st.spinner("Agents collaborating..."):
            result = run_research(topic.strip(), use_memory)

        if result and result.get("report"):
            st.success("Report generated successfully!")
            tabs = st.tabs(["📄 Report", "🔍 Research Notes", "⚖️ Critique", "🧩 Synthesis"])
            with tabs[0]:
                st.markdown(result["report"])
            with tabs[1]:
                st.markdown(result.get("research_notes", ""))
            with tabs[2]:
                st.markdown(result.get("critique", ""))
            with tabs[3]:
                st.markdown(result.get("synthesis", ""))
            if use_memory:
                memory.store(topic.strip(), result["report"])
                st.toast("Saved to cross-session memory", icon="💾")
            st.download_button(
                "Download Report",
                data=result["report"],
                file_name=f"research_{topic[:40].replace(' ', '_')}.md",
                mime="text/markdown",
            )
        else:
            st.error("Research pipeline did not produce a report. Check your API key and try again.")
        st.session_state.running = False

    with st.expander("How it works"):
        st.markdown(
            """
            | Agent | Role |
            |-------|------|
            | **Researcher** | Web search + knowledge synthesis into structured notes |
            | **Critic** | Quality gate — gaps, bias, evidence checks; loops back if needed |
            | **Synthesizer** | Merges research + critique into unified analytical framework |
            | **Writer** | Produces publication-quality markdown report |
            """
        )


if __name__ == "__main__":
    main()