src/research/deep_researcher.py

"""
DeepResearcher — Odysseus IterResearch engine adapted for Job Automation Agent.

Architecture (from Odysseus src/deep_research.py):
  Each round: Think → Search → Extract → Synthesize → Decide (stop/continue)
  Final round: Write polished long-form report.

Adaptations:
  - LLM: uses NVIDIA API (OpenAI-compatible) via httpx directly
  - Search: uses local DuckDuckGo wrapper (no SearXNG needed)
  - No Odysseus internal imports
"""

import asyncio
import json
import logging
import re
import time
from datetime import datetime
from typing import Callable, Dict, List, Optional, Set

from .utils import strip_thinking, is_low_quality, EXTRACTOR_PROMPT
from .search import web_search, fetch_page_content

logger = logging.getLogger(__name__)


def current_date_context() -> str:
    now = datetime.now().astimezone()
    return (
        f"Today's date is {now.strftime('%B %d, %Y')} ({now.strftime('%Y-%m-%d')}). "
        f"When a search query needs a year or refers to 'latest'/'current'/'this year', "
        f"use {now.strftime('%Y')} — never a year inferred from training data.\n\n"
    )


# ── Prompts (verbatim from Odysseus deep_research.py) ──────────────────────
RESEARCH_PLAN_PROMPT = """\
You are a research strategist. Before searching, analyze this question and create a research plan.

**Question:** {question}

Break this question down:
1. What are the key sub-topics that need to be covered for a comprehensive answer?
2. What specific data points, facts, or perspectives should we look for?
3. What would a complete, high-quality answer include?

Return a JSON object with:
- "sub_questions": Array of 3-6 specific sub-questions to investigate
- "key_topics": Array of key topics/angles to cover
- "success_criteria": One sentence describing what a complete answer looks like
"""

QUERY_GEN_PROMPT = """\
You are a research assistant planning web searches.

**Original question:** {question}
**Research plan:** {research_plan}
**What we know so far:** {report}
**Round:** {round_num}

Generate {num_queries} focused search queries that will help answer the question.
{round_instruction}

Return ONLY a JSON array of query strings, nothing else.
Example: ["query one", "query two", "query three"]
"""

SYNTHESIZE_PROMPT = """\
You are updating an evolving research report.

**Original question:** {question}
**Current report:** {report}
**New findings from this round:** {new_findings}

Integrate the new findings into the existing report. Produce an updated, well-organized
report that answers the original question as completely as possible given all evidence.
Remove redundancy, resolve contradictions, maintain logical flow. Keep source URLs as inline citations.

Write only the updated report — no preamble or meta-commentary.
"""

STOP_PROMPT = """\
You are deciding whether a research report is comprehensive enough.

**Original question:** {question}
**Current report:** {report}
**Rounds completed:** {round_num}

Do we have enough information to answer the question comprehensively?
Consider: key aspects addressed? obvious gaps? evidence from multiple sources?

Reply with ONLY "YES" or "NO" followed by a brief one-sentence reason.
Example: "YES — The report covers all major aspects with evidence from multiple sources."
"""

FINAL_REPORT_PROMPT = """\
Write a **detailed, comprehensive** research report answering this question:

**Question:** {question}

**All collected evidence and analysis:**
{report}

Requirements:
- Write at MINIMUM 800 words
- Use clear ## headings and ### subheadings
- Synthesize and analyze — explain WHY things matter
- Include specific data points, numbers, statistics from the evidence
- Include source URLs as inline citations [like this](url)
- Add a brief executive summary at the top
- End with a clear conclusion that directly answers the question
"""

CATEGORY_PROMPTS = {
    "product": "Structure as a RANKED LIST with Pros/Cons per item, quick-compare table, and a Verdict section.",
    "comparison": "Create a Comparison Table, a section per option with strengths/weaknesses, and Best For verdicts.",
    "howto": "Start with a Quick Guide (numbered steps), then Prerequisites, then detailed step sections, then Common Mistakes.",
    "factcheck": "Structure as: The Claim → Evidence For → Evidence Against → Verdict → Nuance & Caveats.",
}


class DeepResearcher:
    """
    Iterative research engine (Odysseus IterResearch pattern).
    Uses DuckDuckGo for search and the NVIDIA API for LLM calls.
    """

    def __init__(
        self,
        llm_endpoint: str,
        llm_model: str,
        llm_api_key: str,
        max_rounds: int = 5,
        max_time: int = 300,
        max_urls_per_round: int = 4,
        max_content_chars: int = 12000,
        max_report_tokens: int = 4096,
        extraction_concurrency: int = 3,
        min_rounds: int = 2,
        progress_callback: Optional[Callable] = None,
        category: Optional[str] = None,
    ):
        self.llm_endpoint     = llm_endpoint
        self.llm_model        = llm_model
        self.llm_api_key      = llm_api_key
        self.max_rounds       = max_rounds
        self.max_time         = max_time
        self.max_urls_per_round = max_urls_per_round
        self.max_content_chars  = max_content_chars
        self.max_report_tokens  = max_report_tokens
        self.extraction_concurrency = extraction_concurrency
        self.min_rounds       = min_rounds
        self._progress        = progress_callback
        self.category         = category
        self._cancelled       = False
        self._start_time      = 0.0
        self.queries_used: Set[str]  = set()
        self.urls_fetched: Set[str]  = set()
        self.round_count      = 0
        self.providers_used: List[str] = []
        self.findings: List[Dict]    = []
        self.evolving_report  = ""
        self.research_plan    = ""

    def cancel(self):
        self._cancelled = True

    # ── Public API ────────────────────────────────────────────────────────
    async def research(self, question: str, prior_report: str = "") -> str:
        self._start_time = time.time()
        findings: List[Dict] = []
        report = prior_report or ""

        self._emit(phase="planning")
        self.research_plan = await self._create_plan(question)  # 120s timeout inside
        if not self.category:
            self.category = await self._classify_category(question)

        consecutive_empty = 0

        for round_num in range(1, self.max_rounds + 1):
            self.round_count = round_num
            if self._cancelled or self._time_exceeded():
                break

            logger.info(f"=== Research Round {round_num} ===")
            self._emit(phase="searching", round=round_num, total_sources=len(self.urls_fetched))

            queries = await self._generate_queries(question, report, round_num)
            if not queries:
                break

            self._emit(phase="searching", round=round_num, queries=len(queries),
                       query_preview=queries[0], total_sources=len(self.urls_fetched))

            round_findings = await self._search_and_extract(queries, question)
            if round_findings:
                findings.extend(round_findings)
                consecutive_empty = 0
                self._emit(phase="reading", round=round_num,
                           new_sources=len(round_findings), total_sources=len(self.urls_fetched))
            else:
                consecutive_empty += 1
                if consecutive_empty >= 2:
                    logger.warning("Search returned nothing for 2 rounds — stopping")
                    break

            if findings:
                self._emit(phase="analyzing", round=round_num)
                report = await self._synthesize(question, findings, report)

            if round_num >= self.min_rounds and await self._should_stop(question, report, round_num):
                logger.info(f"LLM decided to stop after round {round_num}")
                break

        self._emit(phase="writing", total_sources=len(self.urls_fetched))

        if not report:
            if findings:
                return self._fallback_report(question, findings)
            return "No information could be gathered for this question."

        final = await self._final_report(question, report)
        elapsed = time.time() - self._start_time
        logger.info(f"Research complete: {self.round_count} rounds, {len(findings)} findings, {elapsed:.1f}s")
        return final

    # ── LLM helper ────────────────────────────────────────────────────────
    async def _llm(self, messages: List[Dict], temperature: float = 0.3,
                   max_tokens: int = 2048, timeout: int = 300) -> str:
        """
        Calls the LLM using the existing OpenAI SDK client (handles long timeouts).
        Runs the sync call in a thread so the async loop stays free.
        """
        from openai import OpenAI
        def _sync_call() -> str:
            client = OpenAI(base_url="https://integrate.api.nvidia.com/v1",
                            api_key=self.llm_api_key,
                            timeout=timeout)
            resp = client.chat.completions.create(
                model=self.llm_model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=False,
            )
            return resp.choices[0].message.content or ""

        text = await asyncio.to_thread(_sync_call)
        return strip_thinking(text)

    # ── Plan ──────────────────────────────────────────────────────────────
    async def _create_plan(self, question: str) -> str:
        prompt = current_date_context() + RESEARCH_PLAN_PROMPT.format(question=question)
        try:
            response = await self._llm([{"role": "user", "content": prompt}],
                                       max_tokens=512, timeout=180)
            parsed = self._parse_json_object(response)
            if parsed:
                parts = []
                if parsed.get("sub_questions"):
                    parts.append("Sub-questions: " + "; ".join(parsed["sub_questions"]))
                if parsed.get("key_topics"):
                    parts.append("Key topics: " + ", ".join(parsed["key_topics"]))
                return "\n".join(parts) if parts else response
            return response
        except Exception as e:
            logger.warning(f"Planning failed: {e}")
            return ""

    async def _classify_category(self, question: str) -> Optional[str]:
        valid = ", ".join(CATEGORY_PROMPTS.keys())
        prompt = (f"Classify into ONE category: {valid}\n"
                  f"Question: {question}\nRespond with ONLY the category name.")
        try:
            result = await self._llm([{"role": "user", "content": prompt}],
                                     temperature=0, max_tokens=20, timeout=120)
            cat = (result or "").strip().lower().split()[0].strip(".,\"'")
            return cat if cat in CATEGORY_PROMPTS else None
        except Exception:
            return None

    # ── Query generation ─────────────────────────────────────────────────
    async def _generate_queries(self, question: str, report: str,
                                round_num: int) -> List[str]:
        if round_num == 1:
            num_queries, round_instruction = 4, "Generate broad, diverse queries covering key facets."
        else:
            num_queries, round_instruction = 3, "Generate targeted follow-up queries to fill gaps."

        prompt = current_date_context() + QUERY_GEN_PROMPT.format(
            question=question,
            research_plan=self.research_plan or "(No plan — search broadly.)",
            report=report or "(No findings yet.)",
            round_num=round_num,
            num_queries=num_queries,
            round_instruction=round_instruction,
        )
        try:
            response = await self._llm([{"role": "user", "content": prompt}],
                                       temperature=0.5, max_tokens=512, timeout=180)
            queries = self._parse_json_array(response)
            new_queries = [q for q in queries if q not in self.queries_used]
            self.queries_used.update(new_queries)
            return new_queries
        except Exception as e:
            logger.error(f"Query generation failed: {e}")
            return []

    # ── Search + Extract ──────────────────────────────────────────────────
    async def _search_and_extract(self, queries: List[str],
                                  question: str) -> List[Dict]:
        all_findings: List[Dict] = []

        search_tasks = [asyncio.to_thread(web_search, q, 6) for q in queries]
        search_results = await asyncio.gather(*search_tasks, return_exceptions=True)

        urls_to_fetch = []
        for result in search_results:
            if isinstance(result, Exception):
                continue
            for r in (result or []):
                url = r.get("url", "")
                if url and url not in self.urls_fetched:
                    urls_to_fetch.append(r)
                    self.urls_fetched.add(url)
            if len(urls_to_fetch) >= self.max_urls_per_round * len(queries):
                break

        if self._cancelled or self._time_exceeded():
            return all_findings

        semaphore = asyncio.Semaphore(self.extraction_concurrency)

        async def _bounded_extract(r: Dict) -> Optional[Dict]:
            async with semaphore:
                return await self._fetch_and_extract(r["url"], question, r.get("title", ""))

        extract_tasks = [_bounded_extract(r) for r in urls_to_fetch]
        results = await asyncio.gather(*extract_tasks, return_exceptions=True)

        for res in results:
            if isinstance(res, Exception):
                continue
            if res:
                all_findings.append(res)

        return all_findings

    async def _fetch_and_extract(self, url: str, question: str,
                                 title: str) -> Optional[Dict]:
        self._emit(phase="reading", url=url, title=title or url)
        try:
            page = await asyncio.to_thread(fetch_page_content, url, 10)
        except Exception as e:
            logger.warning(f"Fetch failed {url}: {e}")
            return None

        if not page.get("success") or not page.get("content"):
            return None

        content = page["content"]
        if len(content) > self.max_content_chars:
            truncated = content[:self.max_content_chars]
            last_para = truncated.rfind("\n\n")
            content = truncated[:last_para] if last_para > self.max_content_chars * 0.8 else truncated

        prompt = EXTRACTOR_PROMPT.format(webpage_content=content, goal=question)
        try:
            response = await self._llm([{"role": "user", "content": prompt}],
                                       temperature=0.2, max_tokens=1024, timeout=180)
            parsed = self._parse_json_object(response)
            if parsed:
                parsed["url"]   = url
                parsed["title"] = title or page.get("title", "")
                if is_low_quality(parsed.get("summary", "")):
                    return None
                return parsed
            return {
                "url": url, "title": title or page.get("title", ""),
                "rational": "raw", "evidence": response[:2000],
                "summary": response[:400],
            }
        except Exception as e:
            logger.warning(f"LLM extraction failed {url}: {e}")
            # Fallback: use raw page content snippet without LLM
            snippet = content[:600].replace("\n", " ").strip()
            if len(snippet) > 80:
                return {
                    "url": url,
                    "title": title or page.get("title", ""),
                    "rational": "snippet fallback (LLM unavailable)",
                    "evidence": snippet,
                    "summary": snippet[:200],
                }
            return None

    # ── Synthesize ────────────────────────────────────────────────────────
    async def _synthesize(self, question: str, findings: List[Dict],
                          current_report: str) -> str:
        window = findings[-10:]
        findings_text = self._format_findings(window)
        prompt = SYNTHESIZE_PROMPT.format(
            question=question,
            report=current_report or "(First round — no report yet.)",
            new_findings=findings_text,
        )
        try:
            return await self._llm([{"role": "user", "content": prompt}],
                                   temperature=0.3, max_tokens=self.max_report_tokens, timeout=300)
        except Exception as e:
            logger.error(f"Synthesis failed: {e}")
            return current_report

    # ── Stop decision ─────────────────────────────────────────────────────
    async def _should_stop(self, question: str, report: str, round_num: int) -> bool:
        prompt = STOP_PROMPT.format(question=question, report=report, round_num=round_num)
        try:
            response = await self._llm([{"role": "user", "content": prompt}],
                                       temperature=0.1, max_tokens=100)
            clean = strip_thinking(response).strip()
            answer = re.sub(r'^[\s*_`"\'>#\-]+', '', clean).upper()
            return answer.startswith("YES")
        except Exception:
            return False

    # ── Final report ──────────────────────────────────────────────────────
    async def _final_report(self, question: str, report: str) -> str:
        prompt = FINAL_REPORT_PROMPT.format(question=question, report=report)
        cat_extra = CATEGORY_PROMPTS.get(self.category or "", "")
        if cat_extra:
            prompt += f"\n\n**Format note:** {cat_extra}"
        try:
            result = await self._llm([{"role": "user", "content": prompt}],
                                     temperature=0.3, max_tokens=self.max_report_tokens, timeout=180)
            if len(result.split()) < 300:
                expanded = await self._llm(
                    [{"role": "user", "content": prompt},
                     {"role": "assistant", "content": result},
                     {"role": "user", "content": "This is too short. Please expand significantly with more detail, data, and analysis. Target 800+ words."}],
                    temperature=0.4, max_tokens=self.max_report_tokens, timeout=180,
                )
                if len(expanded.split()) > len(result.split()):
                    return expanded
            return result
        except Exception as e:
            logger.error(f"Final report failed: {e}")
            return report

    # ── Helpers ───────────────────────────────────────────────────────────
    def _emit(self, **kwargs):
        if self._progress:
            try:
                self._progress(kwargs)
            except Exception:
                pass

    def _time_exceeded(self) -> bool:
        return (time.time() - self._start_time) > self.max_time

    def _format_findings(self, findings: List[Dict]) -> str:
        parts = []
        for i, f in enumerate(findings, 1):
            url     = f.get("url", "unknown")
            title   = f.get("title", "")
            summary = f.get("summary", "")
            evidence = f.get("evidence", "")
            content = summary if summary else evidence[:800]
            parts.append(f"**Finding {i}** — [{title}]({url})\n{content}")
        return "\n\n".join(parts)

    def _fallback_report(self, question: str, findings: List[Dict]) -> str:
        return (
            f"# {question}\n\n"
            f"_Synthesis did not complete. {len(findings)} finding(s) gathered:_\n\n"
            f"{self._format_findings(findings)}"
        )

    @staticmethod
    def _strip_code_block(text: str) -> str:
        text = text.strip()
        if text.startswith("```"):
            text = re.sub(r'^```(?:json)?\s*', '', text)
            text = re.sub(r'\s*```$', '', text)
        return text.strip()

    def _parse_json_array(self, text: str) -> List[str]:
        text = self._strip_code_block(text)
        try:
            parsed = json.loads(text)
            if isinstance(parsed, list):
                return [str(i) for i in parsed]
        except json.JSONDecodeError:
            pass
        match = re.search(r'\[[\s\S]*\]', text)
        if match:
            try:
                parsed = json.loads(match.group())
                if isinstance(parsed, list):
                    return [str(i) for i in parsed]
            except json.JSONDecodeError:
                pass
        # Last resort: harvest quoted strings
        items = re.findall(r'"([^"]{3,})"', text)
        return items if items else []

    def _parse_json_object(self, text: str) -> Optional[Dict]:
        text = self._strip_code_block(text)
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass
        match = re.search(r'\{[\s\S]*\}', text)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                pass
        return None

    def get_stats(self) -> Dict:
        elapsed = time.time() - self._start_time if self._start_time else 0
        return {
            "Duration":  f"{elapsed:.1f}s",
            "Rounds":    self.round_count,
            "Queries":   len(self.queries_used),
            "URLs":      len(self.urls_fetched),
            "Model":     self.llm_model,
        }


# ── Convenience wrapper ───────────────────────────────────────────────────────
def research_question(question: str, progress_cb: Optional[Callable] = None) -> str:
    """
    Synchronous wrapper — research any question and return a Markdown report.
    Uses GLM 5.1 by default (most reliable in tests).
    """
    import os
    from dotenv import load_dotenv
    load_dotenv()

    api_key   = os.getenv("NVIDIA_API_KEY")
    endpoint  = "https://integrate.api.nvidia.com/v1/chat/completions"
    model     = "z-ai/glm-5.1"

    researcher = DeepResearcher(
        llm_endpoint=endpoint,
        llm_model=model,
        llm_api_key=api_key,
        max_rounds=4,
        max_time=240,
        progress_callback=progress_cb,
    )

    async def _run():
        return await researcher.research(question)

    return asyncio.run(_run())