"""
LangGraph agent-based research strategy with parallel subagent support.

Uses LangChain's create_agent() to build a tool-calling agent that autonomously
decides what to search, when to dig deeper, and when to synthesize. Complex
questions can be decomposed into subtopics researched in parallel by subagents.
"""

from __future__ import annotations

import threading
from concurrent.futures import ThreadPoolExecutor, TimeoutError, as_completed
from datetime import UTC, datetime
from typing import Any, Dict

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langgraph.errors import GraphRecursionError
from loguru import logger

from ...citation_handler import CitationHandler
from ...utilities.search_utilities import (
    extract_links_from_search_results,
    format_links_to_markdown,
)
from .base_strategy import BaseSearchStrategy

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DEFAULT_MAX_ITERATIONS = (
    15  # agent needs many more cycles than pipeline strategies
)
MIN_ITERATIONS = 10  # below this the agent can barely do anything useful
SUBAGENT_TIMEOUT_SECONDS = 300  # 5 minutes per subagent
MAX_SUBTOPICS = 8
MAX_SUBAGENT_WORKERS = 4
CONTENT_FETCH_TIMEOUT = 30
CONTENT_MAX_LENGTH = 10_000


# ---------------------------------------------------------------------------
# Thread-safe search result collector
# ---------------------------------------------------------------------------


class SearchResultsCollector:
    """Accumulates search results from the lead agent and subagents.

    Thread-safe: multiple subagent threads may call ``add_results``
    concurrently.  The ``_all_links`` reference points to the strategy's
    shared ``all_links_of_system`` list and is never reassigned.
    """

    def __init__(self, all_links: list | None = None) -> None:
        self._results: list[dict] = []
        self._sources: list[str] = []
        self._lock = threading.Lock()
        self._all_links = all_links if all_links is not None else []

    # -- public API ----------------------------------------------------------

    def add_results(
        self,
        results: list[dict],
        engine_name: str = "web",
    ) -> int:
        """Index *results* and append to the internal list **and** the shared
        ``all_links_of_system``.  Returns the starting citation index
        (0-based) assigned to the first result in this batch.

        The entire operation runs under a single lock acquisition so that
        citation indices are never duplicated.
        """
        if not results:
            return len(self._all_links)

        with self._lock:
            # Use global offset (all_links) not per-call offset (results)
            # so that indices are unique across sections in detailed reports.
            start_idx = len(self._all_links)
            for i, raw in enumerate(results):
                if not isinstance(raw, dict):
                    continue
                r = dict(raw)  # shallow copy to avoid mutating engine output
                r["index"] = str(start_idx + i + 1)
                r["source_engine"] = engine_name
                # Normalise URL key — citation handler expects "link"
                if "link" not in r and "url" in r:
                    r["link"] = r["url"]
                self._results.append(r)
                link = r.get("link", "")
                if link:
                    self._sources.append(link)
                self._all_links.append(r)
            return start_idx

    def reset(self) -> None:
        """Clear per-call state.  ``_all_links`` is intentionally kept."""
        with self._lock:
            self._results.clear()
            self._sources.clear()

    @property
    def results(self) -> list[dict]:
        with self._lock:
            return list(self._results)

    @property
    def sources(self) -> list[str]:
        with self._lock:
            return list(self._sources)


# ---------------------------------------------------------------------------
# Tool factory helpers
# ---------------------------------------------------------------------------


def _format_results(results: list[dict], start_idx: int) -> str:
    """Format search results as ``[N] Title (URL)\\nSnippet``."""
    lines = []
    for i, r in enumerate(results):
        if not isinstance(r, dict):
            continue
        idx = start_idx + i + 1
        title = r.get("title", "No title")
        link = r.get("link", r.get("url", ""))
        snippet = r.get("snippet", r.get("body", ""))
        lines.append(f"[{idx}] {title} ({link})\n{snippet}")
    return "\n\n".join(lines) if lines else "No results."


def _make_web_search_tool(
    search_engine_name: str,
    model: BaseChatModel,
    settings_snapshot: dict,
    collector: SearchResultsCollector,
):
    """Create a ``web_search`` tool that instantiates a fresh engine per call."""

    @tool
    def web_search(query: str) -> str:
        """Search the web for current information, facts, or news. Returns search result snippets with source indices."""
        from local_deep_research.utilities.resource_utils import safe_close
        from local_deep_research.web_search_engines.search_engine_factory import (
            create_search_engine,
        )

        engine = create_search_engine(
            engine_name=search_engine_name,
            llm=model,
            settings_snapshot=settings_snapshot,
        )
        if engine is None:
            return f"Failed to create search engine '{search_engine_name}'."
        try:
            results = engine.run(query)
            if not isinstance(results, list) or not results:
                return f"No results found for '{query}'. Try rephrasing."
            start = collector.add_results(
                results, engine_name=search_engine_name
            )
            return _format_results(results, start)
        except Exception as exc:
            logger.exception("web_search tool error")
            return f"Search error: {exc}"
        finally:
            safe_close(engine, "web search engine")

    return web_search


def _make_fetch_content_tool(collector: SearchResultsCollector):
    """Create a ``fetch_content`` tool using ContentFetcher."""

    @tool
    def fetch_content(url: str) -> str:
        """Download and read the full text content from a URL. Use when search snippets aren't detailed enough."""
        from local_deep_research.content_fetcher import ContentFetcher

        try:
            with ContentFetcher(timeout=CONTENT_FETCH_TIMEOUT) as fetcher:
                result = fetcher.fetch(url, max_length=CONTENT_MAX_LENGTH)
                if result.get("status") == "success":
                    title = result.get("title", "")
                    content = result.get("content", "")
                    return f"Title: {title}\nURL: {url}\n\n{content}"
                return f"Failed to fetch {url}: {result.get('error', 'unknown error')}"
        except Exception as exc:
            logger.exception("fetch_content tool error")
            return f"Error fetching {url}: {exc}"

    return fetch_content


def _make_specialized_search_tool(
    engine_name: str,
    description: str,
    model: BaseChatModel,
    settings_snapshot: dict,
    collector: SearchResultsCollector,
):
    """Create a ``search_{engine}`` tool for a specific search engine."""

    @tool
    def specialized_search(query: str) -> str:
        """Search a specialized engine."""  # overridden below
        from local_deep_research.utilities.resource_utils import safe_close
        from local_deep_research.web_search_engines.search_engine_factory import (
            create_search_engine,
        )

        engine = create_search_engine(
            engine_name=engine_name,
            llm=model,
            settings_snapshot=settings_snapshot,
        )
        if engine is None:
            return f"Failed to create {engine_name} engine."
        try:
            results = engine.run(query)
            if not isinstance(results, list) or not results:
                return f"No results from {engine_name} for '{query}'. Try rephrasing."
            start = collector.add_results(results, engine_name=engine_name)
            return _format_results(results, start)
        except Exception as exc:
            logger.exception(f"search_{engine_name} tool error")
            return f"Search error ({engine_name}): {exc}"
        finally:
            safe_close(engine, f"{engine_name} search engine")

    # Override name and description after decoration
    specialized_search.name = f"search_{engine_name}"
    specialized_search.description = description
    return specialized_search


def _make_research_subtopic_tool(
    search_engine_name: str,
    model: BaseChatModel,
    settings_snapshot: dict,
    collector: SearchResultsCollector,
    max_sub_iterations: int,
    progress_callback=None,
):
    """Create the ``research_subtopic`` tool that spawns parallel subagents."""

    @tool
    def research_subtopic(subtopics: list[str]) -> str:
        """Delegate parallel research on multiple subtopics. Each subtopic is
        investigated by a separate agent. Pass 2-5 focused research questions."""
        from langchain.agents import create_agent

        if not subtopics:
            return "No subtopics provided."
        if len(subtopics) > MAX_SUBTOPICS:
            subtopics = subtopics[:MAX_SUBTOPICS]

        # Emit progress for UI
        if progress_callback:
            progress_callback(
                f"Researching {len(subtopics)} subtopics in parallel",
                None,
                {
                    "phase": "sub_research",
                    "type": "milestone",
                    "subtopics": subtopics,
                },
            )

        current_date = datetime.now(UTC).strftime("%Y-%m-%d")
        subagent_prompt = (
            f"You are a focused research assistant. Today's date: {current_date}. "
            "Search thoroughly and return a concise factual summary. "
            "Reference sources by their [N] index numbers. "
            "Do NOT ask clarifying questions — provide your findings directly."
        )

        def run_subagent(topic: str) -> str:
            # Each subagent gets its own tool instances (thread safety)
            sub_web_search = _make_web_search_tool(
                search_engine_name, model, settings_snapshot, collector
            )
            sub_fetch = _make_fetch_content_tool(collector)
            try:
                agent = create_agent(
                    model=model,
                    tools=[sub_web_search, sub_fetch],
                    system_prompt=subagent_prompt,
                )
                result = agent.invoke(
                    {"messages": [{"role": "user", "content": topic}]},
                    {"recursion_limit": max_sub_iterations * 2 + 1},
                )
                messages = result.get("messages", [])
                if messages:
                    last = messages[-1]
                    content = getattr(last, "content", str(last))
                    if content:
                        return content
                return f"No findings for: {topic}"
            except GraphRecursionError:
                return f"Research on '{topic}' reached iteration limit. Partial findings above."
            except Exception as exc:
                logger.exception(f"Subagent failed for: {topic[:80]}")
                return f"Research on '{topic}' failed: {exc}"

        ordered_results: dict[str, str] = {}
        num_workers = min(MAX_SUBAGENT_WORKERS, len(subtopics))
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            futures = {executor.submit(run_subagent, t): t for t in subtopics}
            try:
                for future in as_completed(
                    futures, timeout=SUBAGENT_TIMEOUT_SECONDS
                ):
                    topic = futures[future]
                    try:
                        ordered_results[topic] = future.result(timeout=30)
                    except TimeoutError:
                        logger.warning(f"Subagent timed out for: {topic[:80]}")
                        ordered_results[topic] = (
                            f"Research on '{topic}' timed out."
                        )
                    except Exception as exc:
                        ordered_results[topic] = (
                            f"Research on '{topic}' failed: {exc}"
                        )
            except TimeoutError:
                # as_completed itself timed out — some futures didn't finish
                for future, topic in futures.items():
                    if topic not in ordered_results:
                        logger.warning(
                            f"Subagent timed out (overall): {topic[:80]}"
                        )
                        ordered_results[topic] = (
                            f"Research on '{topic}' timed out after "
                            f"{SUBAGENT_TIMEOUT_SECONDS}s."
                        )

        # Return results in original order
        parts = []
        for topic in subtopics:
            parts.append(
                f"## {topic}\n{ordered_results.get(topic, 'No results')}"
            )
        return "\n\n---\n\n".join(parts)

    return research_subtopic


# ---------------------------------------------------------------------------
# Strategy class
# ---------------------------------------------------------------------------


class LangGraphAgentStrategy(BaseSearchStrategy):
    """Research strategy using LangGraph agents with parallel subagent support.

    The lead agent autonomously decides what to search, when to dig deeper
    (via subagents), and when to synthesize — replacing the manual ReAct loop
    in the MCP strategy.
    """

    def __init__(
        self,
        model: BaseChatModel,
        search,
        citation_handler=None,
        max_iterations: int = 15,
        max_sub_iterations: int = 8,
        include_sub_research: bool = True,
        all_links_of_system: list | None = None,
        settings_snapshot: dict | None = None,
        **kwargs,
    ):
        super().__init__(
            all_links_of_system=all_links_of_system,
            settings_snapshot=settings_snapshot,
            **kwargs,
        )
        self.model = model
        self.search = search
        # search.iterations (typically 1-5) controls pipeline strategies.
        # For an agent, each "iteration" is one LLM→tool round-trip, so we
        # need many more.  Treat any value below the agent minimum as "use
        # default" rather than clamping to a uselessly low number.
        self.max_iterations = (
            int(max_iterations)
            if int(max_iterations) >= MIN_ITERATIONS
            else DEFAULT_MAX_ITERATIONS
        )
        self.max_sub_iterations = int(max_sub_iterations)
        self.include_sub_research = include_sub_research
        self.citation_handler = citation_handler or CitationHandler(
            model,
            handler_type="standard",
            settings_snapshot=settings_snapshot,
        )
        self.collector = SearchResultsCollector(self.all_links_of_system)

        # Derive the search engine name for creating fresh instances
        self._search_engine_name = self._resolve_engine_name()

    def _resolve_engine_name(self) -> str:
        """Best-effort extraction of the configured engine name."""
        # Try settings first
        tool_setting = self.get_setting("search.tool", None)
        if tool_setting and isinstance(tool_setting, str):
            return tool_setting
        # Fall back to class name heuristic
        if self.search is not None and hasattr(self.search, "__class__"):
            name = self.search.__class__.__name__
            return name.replace("SearchEngine", "").lower()
        return "duckduckgo"

    def _get_current_engine_name(self) -> str:
        """Get the name of the currently selected search engine."""
        try:
            if hasattr(self.search, "__class__"):
                return self.search.__class__.__name__.replace(
                    "SearchEngine", ""
                ).lower()
        except Exception:
            logger.debug("Could not extract engine name from class")
        return ""

    def _build_tools(self) -> list:
        """Build the LangChain tool list for the lead agent."""
        tools = []

        # Web search (always present if we have a search engine)
        if self.search is not None:
            tools.append(
                _make_web_search_tool(
                    self._search_engine_name,
                    self.model,
                    self.settings_snapshot,
                    self.collector,
                )
            )

        # Content fetcher
        tools.append(_make_fetch_content_tool(self.collector))

        # Specialized search engines
        try:
            from local_deep_research.web_search_engines.search_engines_config import (
                get_available_engines,
            )

            available = get_available_engines(
                settings_snapshot=self.settings_snapshot,
            )
            current = self._get_current_engine_name()
            for name, config in available.items():
                if name in ("auto", "meta") or name == current:
                    continue
                desc = config.get("description", f"Search using {name}")
                strengths = config.get("strengths", [])
                if strengths:
                    desc += f" Best for: {', '.join(strengths[:2])}."
                tools.append(
                    _make_specialized_search_tool(
                        name,
                        desc,
                        self.model,
                        self.settings_snapshot,
                        self.collector,
                    )
                )
        except Exception:
            logger.warning(
                "Failed to load specialized search engines for agent tools"
            )

        # Subagent research tool
        if self.include_sub_research:
            tools.append(
                _make_research_subtopic_tool(
                    self._search_engine_name,
                    self.model,
                    self.settings_snapshot,
                    self.collector,
                    self.max_sub_iterations,
                    progress_callback=self.progress_callback,
                )
            )

        return tools

    # -- Main entry point ---------------------------------------------------

    def analyze_topic(self, query: str) -> Dict[str, Any]:
        from langchain.agents import create_agent

        logger.info(f"LangGraph agent research: {query[:100]}")

        # Reset collector for fresh subsection call (detailed report mode)
        self.collector.reset()
        nr_of_links = len(self.all_links_of_system)

        self._update_progress(
            f'Starting agent research: "{query[:80]}"',
            5,
            {"phase": "init", "type": "milestone", "query": query[:100]},
        )
        self.check_termination()

        # Build tools
        tools = self._build_tools()
        if not tools:
            return self._error_result("No tools available")

        # Build system prompt
        current_date = datetime.now(UTC).strftime("%Y-%m-%d")
        system_prompt = (
            f"You are a research assistant. Today's date: {current_date}.\n"
            "You MUST search the web before answering — never answer from memory alone.\n\n"
            "Strategy:\n"
            "1. Start with web_search for initial exploration.\n"
            "2. For complex multi-faceted questions, use research_subtopic to "
            "investigate specific aspects in parallel (pass 2-5 focused questions).\n"
            "3. Use fetch_content to read full pages when snippets aren't enough.\n"
            "4. Use search_[engine] tools for domain-specific searches "
            "(search_arxiv for science, search_pubmed for medical, etc.).\n"
            "5. When you have enough information, provide a comprehensive answer "
            "citing sources as [1], [2], etc.\n"
        )

        # Create agent — may fail if model doesn't support tool calling
        try:
            agent = create_agent(
                model=self.model,
                tools=tools,
                system_prompt=system_prompt,
            )
        except Exception as exc:
            logger.exception("Failed to create LangGraph agent")
            return self._error_result(
                f"Failed to create agent (model may not support tool calling): {exc}"
            )

        # Stream agent execution
        effective_max = max(MIN_ITERATIONS, self.max_iterations)
        config = {"recursion_limit": effective_max * 2 + 1}
        iteration = 0
        final_content = ""
        agent_messages: list = []

        try:
            for chunk in agent.stream(
                {"messages": [{"role": "user", "content": query}]},
                config,
                stream_mode="updates",
            ):
                self.check_termination()

                if "agent" in chunk or "model" in chunk:
                    node_key = "agent" if "agent" in chunk else "model"
                    iteration += 1
                    progress = 10 + int((iteration / effective_max) * 75)
                    msgs = chunk[node_key].get("messages", [])
                    for msg in msgs:
                        if isinstance(msg, AIMessage):
                            agent_messages.append(msg)
                            content = msg.content or ""
                            tool_calls = getattr(msg, "tool_calls", [])

                            if tool_calls:
                                for tc in tool_calls:
                                    tc_args = tc.get("args", {})
                                    preview = str(
                                        tc_args.get(
                                            "query", tc_args.get("url", "")
                                        )
                                    )[:80]
                                    self._update_progress(
                                        f'Tool: {tc["name"]} — "{preview}"',
                                        min(85, progress),
                                        {
                                            "phase": "tool_call",
                                            "tool": tc["name"],
                                            "iteration": iteration,
                                        },
                                    )
                            elif content:
                                # No tool calls = final answer
                                final_content = content

                elif "tools" in chunk:
                    msgs = chunk["tools"].get("messages", [])
                    for msg in msgs:
                        tool_name = getattr(msg, "name", "tool")
                        preview = str(getattr(msg, "content", ""))[
                            :150
                        ].replace("\n", " ")
                        self._update_progress(
                            f"Result from {tool_name}: {preview}",
                            min(
                                85,
                                10 + int((iteration / effective_max) * 75) + 3,
                            ),
                            {"phase": "observation", "tool": tool_name},
                        )

        except GraphRecursionError:
            logger.warning(
                "LangGraph agent hit recursion limit, synthesizing partial results"
            )
            if not final_content:
                final_content = self._synthesize_from_collector(query)
        except Exception as exc:
            logger.exception("LangGraph agent error")
            if not final_content:
                if self.collector.results:
                    final_content = self._synthesize_from_collector(query)
                else:
                    return self._error_result(f"Agent error: {exc}")

        if not final_content:
            if self.collector.results:
                final_content = self._synthesize_from_collector(query)
            else:
                final_content = (
                    "Research could not produce results. Try a different query."
                )

        return self._finalize(
            query, final_content, iteration, nr_of_links, agent_messages
        )

    # -- Helpers ------------------------------------------------------------

    def _synthesize_from_collector(self, query: str) -> str:
        """Fallback synthesis when the agent was cut short."""
        results = self.collector.results
        if not results:
            return "Research could not be completed within the iteration limit."
        summaries = []
        for r in results[:20]:
            summaries.append(
                f"[{r.get('index', '?')}] {r.get('title', '')}: "
                f"{r.get('snippet', '')}"
            )
        prompt = (
            f"Synthesize a comprehensive answer to: {query}\n\n"
            f"Based on these sources:\n" + "\n".join(summaries)
        )
        try:
            response = self.model.invoke(prompt)
            return (
                response.content
                if hasattr(response, "content")
                else str(response)
            )
        except Exception as exc:
            logger.exception("Fallback synthesis failed")
            return f"Research collected {len(results)} sources but synthesis failed: {exc}"

    def _finalize(
        self,
        query: str,
        final_answer: str,
        iteration: int,
        nr_of_links: int,
        agent_messages: list,
    ) -> Dict[str, Any]:
        """Apply citation handling and build the return dict."""
        self._update_progress(
            f"Synthesizing {len(self.collector.results)} sources with citations",
            90,
            {"phase": "synthesis", "type": "milestone"},
        )

        all_search_results = self.collector.results
        synthesized_content = final_answer
        documents: list = []

        # Citation handling — only if we have results
        if all_search_results:
            try:
                citation_result = self.citation_handler.analyze_followup(
                    query,
                    all_search_results,
                    previous_knowledge=final_answer,
                    nr_of_links=nr_of_links,
                )
                if isinstance(citation_result, dict):
                    synthesized_content = citation_result.get(
                        "content", citation_result.get("response", final_answer)
                    )
                    documents = citation_result.get("documents", [])
            except Exception:
                logger.warning(
                    "Citation handler failed, using raw agent answer"
                )

        # Format sources
        formatted_output = synthesized_content
        if all_search_results:
            try:
                all_links = extract_links_from_search_results(
                    all_search_results
                )
                if all_links:
                    sources_md = format_links_to_markdown(all_links)
                    if sources_md:
                        formatted_output = f"{synthesized_content}\n\n## Sources\n\n{sources_md}"
            except Exception:
                logger.exception("Failed to format source links")

        # Build reasoning trace from agent messages
        reasoning_trace = []
        for msg in agent_messages:
            entry: Dict[str, Any] = {"role": "assistant"}
            if hasattr(msg, "content") and msg.content:
                entry["content"] = msg.content
            tool_calls = getattr(msg, "tool_calls", [])
            if tool_calls:
                entry["tool_calls"] = [
                    {"name": tc.get("name"), "args": tc.get("args", {})}
                    for tc in tool_calls
                ]
            reasoning_trace.append(entry)

        self._update_progress(
            "Research complete",
            100,
            {"phase": "complete", "type": "milestone", "iterations": iteration},
        )

        return {
            "findings": [
                {
                    "content": synthesized_content,
                    "question": query,
                    "search_results": all_search_results,
                    "documents": documents,
                }
            ],
            "iterations": iteration,
            "questions": {},
            "formatted_findings": formatted_output,
            "current_knowledge": synthesized_content,
            "sources": list(set(self.collector.sources)),
            "search_results": all_search_results,
            "documents": documents,
            "reasoning_trace": reasoning_trace,
            "error": None,
        }

    def _error_result(self, error: str) -> Dict[str, Any]:
        logger.error(f"LangGraph agent strategy error: {error}")
        self._update_progress(
            f"Error: {error}",
            100,
            {"phase": "error", "error": error, "status": "failed"},
        )
        return {
            "findings": [],
            "iterations": 0,
            "questions": {},
            "formatted_findings": f"Error: {error}",
            "current_knowledge": "",
            "sources": [],
            "search_results": [],
            "documents": [],
            "reasoning_trace": [],
            "error": error,
        }

    def close(self):
        """No persistent resources to clean up."""
        pass
