"""Context gatherer for pre-planning exploration.

Before creating a plan, the agent needs to understand the task context:
- What files/code exist in the workspace?
- What documentation is available online?
- What's the current state of the system?

This component uses tools to gather that context.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any

from promethian.config import PromethianConfig
from promethian.core.types import Skill, Tool
from promethian.utils.claude import ClaudeClient

# Import built-in tools
from promethian.tools.bash import get_bash, BashTool
from promethian.tools.web_search import get_web_search, WebSearchTool
from promethian.tools.web_fetch import get_web_fetch, WebFetchTool

# Browser is optional
try:
    from promethian.tools.browser import get_browser, BrowserTool
    BROWSER_AVAILABLE = True
except ImportError:
    BROWSER_AVAILABLE = False
    get_browser = None
    BrowserTool = None


@dataclass
class GatheredContext:
    """Context gathered before planning."""

    # File/codebase context
    files_examined: list[str] = field(default_factory=list)
    file_contents: dict[str, str] = field(default_factory=dict)
    directory_structure: str | None = None

    # Search results
    web_search_results: list[dict] = field(default_factory=list)

    # Command outputs
    command_outputs: dict[str, str] = field(default_factory=dict)

    # Browser data
    page_contents: dict[str, str] = field(default_factory=dict)

    # Summary for the planner
    summary: str = ""

    # Raw tool results for reference
    raw_results: list[dict] = field(default_factory=list)

    @property
    def sources(self) -> list[str]:
        """Return a list of all sources examined (files + URLs)."""
        sources = []
        sources.extend(self.files_examined)
        for result in self.web_search_results:
            if url := result.get("url"):
                sources.append(url)
        sources.extend(self.page_contents.keys())
        return sources

    def to_prompt_context(self) -> str:
        """Convert gathered context to a string for the planner."""
        parts = []

        if self.summary:
            parts.append(f"# Context Summary\n{self.summary}")

        if self.directory_structure:
            parts.append(f"# Project Structure\n```\n{self.directory_structure}\n```")

        if self.file_contents:
            parts.append("# Relevant Files")
            for path, content in self.file_contents.items():
                # Truncate long files
                if len(content) > 2000:
                    content = content[:2000] + "\n... (truncated)"
                parts.append(f"\n## {path}\n```\n{content}\n```")

        if self.web_search_results:
            parts.append("# Web Search Results")
            for result in self.web_search_results[:5]:
                parts.append(f"- [{result.get('title', 'No title')}]({result.get('url', '')})")
                if result.get('snippet'):
                    parts.append(f"  {result['snippet'][:200]}")

        if self.page_contents:
            parts.append("# Website Structure (USE THESE PATHS)")
            for url, content in self.page_contents.items():
                parts.append(f"\n## {url}\n{content}")

        if self.command_outputs:
            parts.append("# Command Outputs")
            for cmd, output in self.command_outputs.items():
                if len(output) > 500:
                    output = output[:500] + "\n... (truncated)"
                parts.append(f"\n## `{cmd}`\n```\n{output}\n```")

        return "\n\n".join(parts) if parts else ""


CONTEXT_ANALYSIS_PROMPT = """You are a context analysis assistant. Your job is to determine what context needs to be gathered before planning a task.

Given a task, analyze what information would be helpful to gather BEFORE creating an execution plan.

Consider:
1. **File/Code Context**: What files or code might be relevant? What directory structure matters?
2. **Documentation**: Would web searches for documentation or examples help?
3. **System State**: What commands could reveal the current state (git status, package versions, etc.)?
4. **External Websites/APIs**: If the task involves a website, organization, or external service, ALWAYS search for the correct URL first!

IMPORTANT: For tasks involving external websites or web scraping:
- ALWAYS use web_search to find the correct website URL
- Website domains vary (.com, .org, .co.uk, .org.uk, etc.) - never guess!
- Search for "[organization name] official website" or "[organization name] [specific page]"

Available tools for gathering context:
- `bash`: Run shell commands (ls, cat, git, find, grep, etc.)
- `web_search`: Search the web for documentation, tutorials, URLs, examples
- `web_news_search`: Search for recent news/updates
- `web_fetch`: Fetch a URL directly (use AFTER finding the URL with web_search)
{browser_tools}

Output a JSON plan for gathering context:
{{
    "reasoning": "Why this context is needed",
    "context_steps": [
        {{
            "tool": "web_search",
            "purpose": "Find the official website URL",
            "params": {{"query": "Acme Corp official website"}}
        }},
        {{
            "tool": "web_fetch",
            "purpose": "Fetch homepage to find navigation links",
            "params": {{"url": "USE_URL_FROM_SEARCH_RESULTS"}}
        }},
        {{
            "tool": "bash",
            "purpose": "List project files",
            "params": {{"command": "ls -la"}}
        }}
    ],
    "max_steps": 5
}}

Keep it focused - only gather context that's directly relevant to the task.
For local coding tasks that don't involve external services, you may return an empty context_steps array.
For tasks involving external websites or services, ALWAYS include a web_search step to find the correct URL."""


class ContextGatherer:
    """
    Gathers context before planning using available tools.

    This enables the agent to:
    - Explore the codebase before planning
    - Search for documentation
    - Check system state
    - Understand the environment

    Uses a fast model (Haiku by default) for context gathering
    since it's mostly analysis and tool coordination.
    """

    def __init__(
        self,
        config: PromethianConfig,
        claude: ClaudeClient | None = None,
        enable_bash: bool = True,
        enable_web_search: bool = True,
        enable_browser: bool = True,
        max_context_steps: int = 10,
    ):
        self.config = config
        self.max_context_steps = max_context_steps

        # Use a separate Claude client with context-gathering model (Haiku by default)
        if claude is None:
            self.claude = ClaudeClient(
                api_key=config.anthropic_api_key,
                model=config.get_context_model(),
                max_tokens=2048,  # Context analysis needs less tokens
            )
        else:
            self.claude = claude

        # Tool availability
        self._enable_bash = enable_bash
        self._enable_web_search = enable_web_search
        self._enable_browser = enable_browser and BROWSER_AVAILABLE

        # Lazy tool instances
        self._bash: BashTool | None = None
        self._web_search: WebSearchTool | None = None
        self._web_fetch: WebFetchTool | None = None
        self._browser: BrowserTool | None = None

    def _get_bash(self) -> BashTool:
        if self._bash is None:
            self._bash = get_bash(
                timeout=self.config.tool_timeout,
                working_dir=self.config.working_dir,
            )
        return self._bash

    def _get_web_search(self) -> WebSearchTool:
        if self._web_search is None:
            self._web_search = get_web_search()
        return self._web_search

    def _get_web_fetch(self) -> WebFetchTool:
        if self._web_fetch is None:
            self._web_fetch = get_web_fetch()
        return self._web_fetch

    async def _get_browser(self) -> BrowserTool:
        if self._browser is None and self._enable_browser:
            self._browser = await get_browser()
        return self._browser

    async def gather_context(
        self,
        task: str,
        skills: list[Skill] | None = None,
        tools: list[Tool] | None = None,
    ) -> GatheredContext:
        """
        Gather context needed for planning a task.

        Args:
            task: The task to gather context for
            skills: Available skills (for reference)
            tools: Available tools (custom tools that could be used)

        Returns:
            GatheredContext with all gathered information
        """
        # Analyze what context is needed
        context_plan = await self._analyze_context_needs(task, skills, tools)

        # Execute context gathering steps
        context = GatheredContext()

        steps = context_plan.get("context_steps", [])
        max_steps = min(
            context_plan.get("max_steps", self.max_context_steps),
            self.max_context_steps
        )

        for step in steps[:max_steps]:
            tool_name = step.get("tool")
            params = step.get("params", {})
            purpose = step.get("purpose", "")

            result = await self._execute_context_step(tool_name, params)

            # Store result
            context.raw_results.append({
                "tool": tool_name,
                "purpose": purpose,
                "params": params,
                "result": result,
            })

            # Categorize result
            self._categorize_result(context, tool_name, params, result)

            # Auto-fetch homepage URLs from search results
            if tool_name in ("web_search", "web_news_search") and result.get("success"):
                await self._auto_fetch_homepages(result, context)

        # Generate summary
        if context.raw_results:
            context.summary = await self._summarize_context(task, context)

        return context

    async def _auto_fetch_homepages(
        self,
        search_result: dict,
        context: GatheredContext,
    ) -> None:
        """Automatically fetch homepage URLs found in search results."""
        from urllib.parse import urlparse

        results = search_result.get("results", [])
        fetched_domains = set()

        for r in results[:3]:  # Check first 3 results
            url = r.get("url", "")
            if not url:
                continue

            parsed = urlparse(url)
            domain = parsed.netloc

            # Skip if we've already fetched this domain
            if domain in fetched_domains:
                continue

            # Check if this looks like a homepage (short path)
            path = parsed.path.strip("/")
            if path and len(path.split("/")) > 1:
                continue  # Skip deep URLs

            # Fetch the homepage
            fetch_result = await self._execute_context_step("web_fetch", {"url": url})

            if fetch_result.get("success"):
                context.raw_results.append({
                    "tool": "web_fetch",
                    "purpose": f"Auto-fetch homepage from search results",
                    "params": {"url": url},
                    "result": fetch_result,
                })
                self._categorize_result(context, "web_fetch", {"url": url}, fetch_result)
                fetched_domains.add(domain)

    async def _analyze_context_needs(
        self,
        task: str,
        skills: list[Skill] | None,
        tools: list[Tool] | None,
    ) -> dict:
        """Analyze what context needs to be gathered."""
        # Build available tools description
        browser_tools = ""
        if self._enable_browser:
            browser_tools = "- `browser_navigate`: Navigate to a URL and get page content"

        system = CONTEXT_ANALYSIS_PROMPT.format(browser_tools=browser_tools)

        # Include info about available skills/tools
        prompt_parts = [f"# Task\n{task}"]

        if skills:
            prompt_parts.append("\n# Available Skills (for reference)")
            for s in skills[:5]:
                prompt_parts.append(f"- {s.name}: {s.description}")

        if tools:
            prompt_parts.append("\n# Available Custom Tools")
            for t in tools[:5]:
                prompt_parts.append(f"- {t.name}: {t.description}")

        prompt_parts.append("\nWhat context should be gathered before planning this task?")
        prompt = "\n".join(prompt_parts)

        try:
            return await self.claude.complete_json(
                prompt=prompt,
                system=system,
            )
        except Exception:
            # If analysis fails, return empty plan
            return {"context_steps": [], "reasoning": "Analysis failed"}

    async def _execute_context_step(
        self,
        tool_name: str,
        params: dict,
    ) -> dict[str, Any]:
        """Execute a single context gathering step."""
        try:
            if tool_name == "bash" and self._enable_bash:
                bash = self._get_bash()
                result = await bash.run(
                    command=params.get("command", "echo 'No command'"),
                    timeout=30,  # Short timeout for context gathering
                )
                return {
                    "success": result.success,
                    "stdout": result.stdout,
                    "stderr": result.stderr,
                    "exit_code": result.exit_code,
                }

            elif tool_name == "web_search" and self._enable_web_search:
                ws = self._get_web_search()
                result = await ws.search(
                    query=params.get("query", ""),
                    max_results=params.get("max_results", 5),
                )
                if result.success:
                    return {
                        "success": True,
                        "results": [
                            {"title": r.title, "url": r.url, "snippet": r.snippet}
                            for r in (result.results or [])
                        ],
                    }
                return {"success": False, "error": result.error}

            elif tool_name == "web_news_search" and self._enable_web_search:
                ws = self._get_web_search()
                result = await ws.news_search(
                    query=params.get("query", ""),
                    max_results=params.get("max_results", 5),
                )
                if result.success:
                    return {
                        "success": True,
                        "results": [
                            {"title": r.title, "url": r.url, "snippet": r.snippet}
                            for r in (result.results or [])
                        ],
                    }
                return {"success": False, "error": result.error}

            elif tool_name == "web_fetch":
                wf = self._get_web_fetch()
                result = await wf.fetch(params.get("url", ""))
                if result.success:
                    return {
                        "success": True,
                        "url": result.url,
                        "content": result.content,
                        "content_type": result.content_type,
                    }
                return {"success": False, "error": result.error}

            elif tool_name == "browser_navigate" and self._enable_browser:
                browser = await self._get_browser()
                result = await browser.navigate(params.get("url", ""))
                if result.success:
                    # Also get page text
                    text_result = await browser.get_text("body")
                    return {
                        "success": True,
                        "title": result.data.get("title"),
                        "url": result.data.get("url"),
                        "text": text_result.data.get("text", "")[:5000] if text_result.success else "",
                    }
                return {"success": False, "error": result.error}

            else:
                return {"success": False, "error": f"Unknown or disabled tool: {tool_name}"}

        except Exception as e:
            return {"success": False, "error": str(e)}

    def _categorize_result(
        self,
        context: GatheredContext,
        tool_name: str,
        params: dict,
        result: dict,
    ) -> None:
        """Categorize a result into the appropriate context field."""
        if not result.get("success"):
            return

        if tool_name == "bash":
            command = params.get("command", "")
            stdout = result.get("stdout", "")

            # Categorize by command type
            if command.startswith(("ls", "find", "tree")):
                if not context.directory_structure:
                    context.directory_structure = stdout
                else:
                    context.directory_structure += f"\n\n{stdout}"

            elif command.startswith("cat ") or command.startswith("head ") or command.startswith("tail "):
                # Extract file path from command
                parts = command.split()
                if len(parts) >= 2:
                    file_path = parts[1]
                    context.files_examined.append(file_path)
                    context.file_contents[file_path] = stdout

            else:
                context.command_outputs[command] = stdout

        elif tool_name in ("web_search", "web_news_search"):
            results = result.get("results", [])
            context.web_search_results.extend(results)

        elif tool_name == "browser_navigate":
            url = result.get("url", params.get("url", ""))
            text = result.get("text", "")
            if url and text:
                context.page_contents[url] = text

        elif tool_name == "web_fetch":
            # Extract useful structure from fetched pages
            url = result.get("url", params.get("url", ""))
            content = result.get("content", "")
            if url and content:
                # Extract navigation links from HTML
                links = self._extract_links(content, url)
                if links:
                    context.page_contents[url] = f"Navigation links found:\n{links}"

    def _extract_links(self, html_content: str, base_url: str) -> str:
        """Extract navigation links from HTML content."""
        import re
        from urllib.parse import urljoin, urlparse

        base_domain = urlparse(base_url).netloc
        links = []
        seen = set()

        # Find all href links
        href_pattern = r'href=["\']([^"\']+)["\']'
        for match in re.finditer(href_pattern, html_content, re.IGNORECASE):
            href = match.group(1)
            if not href or href.startswith(('#', 'javascript:', 'mailto:', 'tel:')):
                continue

            # Convert to absolute URL
            full_url = urljoin(base_url, href)
            parsed = urlparse(full_url)

            # Only include same-domain links
            if parsed.netloc != base_domain:
                continue

            # Get the path
            path = parsed.path
            if not path or path == '/':
                continue

            # Skip static assets
            if any(path.endswith(ext) for ext in ['.css', '.js', '.png', '.jpg', '.gif', '.svg', '.ico', '.woff', '.woff2']):
                continue

            # Deduplicate
            if path in seen:
                continue
            seen.add(path)

            links.append(path)

        # Return sorted unique paths
        if links:
            return '\n'.join(f"- {path}" for path in sorted(set(links))[:30])
        return ""

    async def _summarize_context(
        self,
        task: str,
        context: GatheredContext,
    ) -> str:
        """Generate a summary of gathered context."""
        # Build a description of what was gathered
        parts = [f"Task: {task}", "\nGathered information:"]

        for result in context.raw_results:
            if result["result"].get("success"):
                parts.append(f"- {result['purpose']}: Success")
            else:
                parts.append(f"- {result['purpose']}: Failed - {result['result'].get('error', 'Unknown error')}")

        prompt = "\n".join(parts) + "\n\nProvide a brief (2-3 sentence) summary of the relevant context for planning this task."

        try:
            return await self.claude.complete(prompt=prompt)
        except Exception:
            return "Context gathered but summary generation failed."

    async def close(self) -> None:
        """Clean up resources."""
        if self._browser is not None:
            await self._browser.close()
            self._browser = None
