"""
automl_lib.research
~~~~~~~~~~~~~~~~~~~

WebResearch Engine for automated scraping, analysis, and LLM-based answering.
"""

from __future__ import annotations

import asyncio
import re
from typing import Any, Dict, List

from .utils import get_logger


class WebResearchEngine:
    """
    A research engine that searches the web, scrapes content, 
    builds a local vector DB, and answers questions using a local LLM via Ollama.
    """

    def __init__(
        self,
        embedding_model_name: str = "all-MiniLM-L6-v2",
        llm_model: str = "llama3",
        llm_provider: str = "ollama",
        hf_token: str | None = None,
        max_sources: int = 5,
        scraping_depth: int = 1,
        style: str = "concise",
        research_mode: str = "general",
    ):
        self.logger = get_logger("automl_lib.research")
        self.embedding_model_name = embedding_model_name
        self.llm_model = llm_model
        self.llm_provider = llm_provider.lower()
        self.hf_token = hf_token
        self.max_sources = max_sources
        self.scraping_depth = scraping_depth
        self.style = style
        self.research_mode = research_mode

        self._embedding_model = None
        self._llm_pipeline = None
        self._faiss_index = None
        self._scraped_chunks: List[str] = []
        self._scraped_urls: List[str] = []

    def _ensure_dependencies(self):
        """Check if optional dependencies are installed."""
        try:
            import aiohttp
            import bs4
            from ddgs import AsyncDDGS
            from sentence_transformers import SentenceTransformer
            import faiss
            import numpy as np
            
            if self.llm_provider == "ollama":
                import ollama
            elif self.llm_provider == "huggingface":
                from transformers import pipeline
        except ImportError as e:
            raise ImportError(
                "WebResearch Engine requires additional dependencies. "
                "Install them via: pip install trainiq[research]\n"
                f"Original error: {e}"
            )

    def _init_models(self):
        """Initialize local ML models (embeddings and LLM)."""
        if self._embedding_model is None:
            self.logger.info(f"Loading embedding model: {self.embedding_model_name}")
            from sentence_transformers import SentenceTransformer
            self._embedding_model = SentenceTransformer(self.embedding_model_name)
        
        if self.llm_provider == "huggingface" and self._llm_pipeline is None:
            self.logger.info(f"Loading HuggingFace LLM: {self.llm_model}")
            from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
            import torch
            
            device = 0 if torch.cuda.is_available() else -1
            
            # Detect if model is T5-based (seq2seq) or causal LM
            try:
                tokenizer = AutoTokenizer.from_pretrained(self.llm_model, token=self.hf_token)
                
                # Try loading as seq2seq first (for T5, BART, etc.)
                try:
                    model = AutoModelForSeq2SeqLM.from_pretrained(self.llm_model, token=self.hf_token)
                    task = "text2text-generation"
                except:
                    # Fall back to causal LM (for GPT, Llama, etc.)
                    model = AutoModelForCausalLM.from_pretrained(self.llm_model, token=self.hf_token)
                    task = "text-generation"
                
                self._llm_pipeline = pipeline(
                    task,
                    model=model,
                    tokenizer=tokenizer,
                    device=device,
                    max_new_tokens=512,
                    do_sample=True,
                    temperature=0.7,
                )
            except Exception as e:
                self.logger.error(f"Failed to load HuggingFace model: {e}")
                raise

    async def _search_web(self, query: str) -> List[Dict[str, str]]:
        """Asynchronously search duckduckgo for the query."""
        self.logger.info(f"Searching web for: {query}")
        from ddgs import AsyncDDGS

        # Adjust query based on research mode
        search_query = query
        if self.research_mode == "academic":
            search_query += " site:arxiv.org OR site:edu"
        elif self.research_mode == "github":
            search_query += " site:github.com"
        elif self.research_mode == "documentation":
            search_query += " documentation OR docs"

        results = []
        async with AsyncDDGS() as ddgs:
            try:
                # Use duckduckgo text search
                async for r in ddgs.text(search_query, max_results=self.max_sources):
                    results.append(r)
            except Exception as e:
                self.logger.error(f"Search failed: {e}")

        # Format is usually [{'title': '...', 'href': '...', 'body': '...'}, ...]
        return results

    async def _scrape_page(self, session: Any, url: str) -> str:
        """Fetch and clean HTML content."""
        from bs4 import BeautifulSoup
        try:
            async with session.get(url, timeout=10) as response:
                if response.status == 200:
                    html = await response.text()
                    soup = BeautifulSoup(html, "html.parser")

                    # Remove unwanted tags
                    for script in soup(["script", "style", "header", "footer", "nav", "aside"]):
                        script.extract()

                    # Extract text from paragraphs
                    paragraphs = soup.find_all("p")
                    text = " ".join([p.get_text() for p in paragraphs])

                    # Clean whitespace
                    text = re.sub(r"\s+", " ", text).strip()
                    return text
                else:
                    self.logger.warning(f"Failed to fetch {url}: Status {response.status}")
        except Exception as e:
            self.logger.warning(f"Error scraping {url}: {e}")
        return ""

    async def _scrape_all(self, search_results: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """Scrape all search results concurrently."""
        import aiohttp
        self.logger.info(f"Scraping {len(search_results)} pages...")
        
        scraped_data = []
        async with aiohttp.ClientSession() as session:
            tasks = []
            for res in search_results:
                url = res.get("href", res.get("link", ""))
                if url:
                    tasks.append(self._scrape_page(session, url))

            contents = await asyncio.gather(*tasks)

            for res, content in zip(search_results, contents):
                url = res.get("href", res.get("link", ""))
                if content:
                    scraped_data.append({"url": url, "title": res.get("title", ""), "content": content})

        return scraped_data

    def _chunk_text(self, text: str, chunk_size: int = 150) -> List[str]:
        """Split text into chunks of `chunk_size` words."""
        words = text.split()
        return [" ".join(words[i : i + chunk_size]) for i in range(0, len(words), chunk_size)]

    def _build_vector_db(self, scraped_data: List[Dict[str, str]]):
        """Create FAISS DB with embeddings for all scraped text chunks."""
        self.logger.info("Building local vector database...")
        import faiss
        import numpy as np

        self._scraped_chunks = []
        self._scraped_urls = []

        for data in scraped_data:
            chunks = self._chunk_text(data["content"])
            for chunk in chunks:
                if len(chunk) > 50:  # Only keep meaningful chunks
                    self._scraped_chunks.append(chunk)
                    self._scraped_urls.append(data["url"])

        if not self._scraped_chunks:
            self.logger.warning("No valid content found to build vector DB.")
            return

        embeddings = self._embedding_model.encode(self._scraped_chunks)
        dim = embeddings.shape[1]

        self._faiss_index = faiss.IndexFlatL2(dim)
        self._faiss_index.add(np.array(embeddings).astype("float32"))
        
        self.logger.info(f"Vector DB built with {len(self._scraped_chunks)} chunks.")

    def _retrieve_context(self, query: str, top_k: int = 3) -> List[Dict[str, str]]:
        """Retrieve the top-K chunks from the DB matching the query."""
        if not self._faiss_index or not self._scraped_chunks:
            return []

        import numpy as np

        query_embedding = self._embedding_model.encode([query])
        distances, indices = self._faiss_index.search(np.array(query_embedding).astype("float32"), top_k)

        retrieved = []
        for idx in indices[0]:
            if idx != -1 and idx < len(self._scraped_chunks):
                retrieved.append({
                    "text": self._scraped_chunks[idx],
                    "url": self._scraped_urls[idx]
                })
        return retrieved

    def _generate_answer(self, query: str, context: List[Dict[str, str]]) -> Dict[str, Any]:
        """Call LLM (Ollama or HuggingFace) to synthesize an answer from the retrieved contexts."""
        self.logger.info(f"Generating answer using {self.llm_model} via {self.llm_provider}...")

        if not context:
            return {
                "query": query,
                "answer": "No relevant context could be found to answer the query.",
                "references": [],
                "context_used": 0
            }

        context_str = ""
        unique_urls = set()
        for i, ctx in enumerate(context):
            context_str += f"[Source {i+1}]: {ctx['text']}\n\n"
            unique_urls.add(ctx['url'])

        system_prompt = (
            f"You are a helpful research assistant. Answer the user's query based ONLY on the provided context. "
            f"Your style should be {self.style}. "
            f"Include extracted key insights and cite your sources."
        )

        user_prompt = (
            f"Context information is below.\n"
            f"---------------------\n{context_str}\n---------------------\n"
            f"Query: {query}\nAnswer:"
        )

        try:
            if self.llm_provider == "ollama":
                import ollama
                response = ollama.chat(
                    model=self.llm_model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                )
                answer = response["message"]["content"]
            
            elif self.llm_provider == "huggingface":
                full_prompt = f"{system_prompt}\n\n{user_prompt}"
                
                # Truncate if too long (max ~2000 chars to be safe)
                if len(full_prompt) > 2000:
                    # Keep system prompt and truncate context
                    context_limit = 1500 - len(system_prompt) - len(query) - 100
                    if len(context_str) > context_limit:
                        context_str = context_str[:context_limit] + "...[truncated]"
                        user_prompt = (
                            f"Context information is below.\n"
                            f"---------------------\n{context_str}\n---------------------\n"
                            f"Query: {query}\nAnswer:"
                        )
                        full_prompt = f"{system_prompt}\n\n{user_prompt}"
                
                response = self._llm_pipeline(full_prompt, truncation=True, max_length=512)
                answer = response[0]["generated_text"]
                # Remove the prompt from the generated text if present
                if full_prompt in answer:
                    answer = answer.replace(full_prompt, "").strip()
            
            else:
                raise ValueError(f"Unsupported LLM provider: {self.llm_provider}")
                
        except Exception as e:
            self.logger.error(f"Failed to generate answer via {self.llm_provider}: {e}")
            if self.llm_provider == "ollama":
                answer = f"Error generating answer: {e}. Ensure Ollama is running and the model '{self.llm_model}' is pulled."
            else:
                answer = f"Error generating answer: {e}. Check your HuggingFace model and token."

        return {
            "query": query,
            "answer": answer,
            "references": list(unique_urls),
            "context_used": len(context),
        }

    async def a_run(self, query: str) -> Dict[str, Any]:
        """Async pipeline execution."""
        self._ensure_dependencies()
        self._init_models()

        # 1. Search Web
        search_results = await self._search_web(query)
        if not search_results:
            return {"query": query, "answer": "No results found on the web.", "references": []}

        # 2. Scrape Pages
        scraped_data = await self._scrape_all(search_results)

        # 3. Build Vector DB
        self._build_vector_db(scraped_data)

        # 4. Retrieve Relevant Context
        context = self._retrieve_context(query, top_k=5)

        # 5. Generate Answer
        result = self._generate_answer(query, context)

        return result

    def run(self, query: str) -> Dict[str, Any]:
        """Synchronous wrapper for the research pipeline."""
        return asyncio.run(self.a_run(query))


def research(query: str, **kwargs) -> Dict[str, Any]:
    """
    Executes a web research pipeline for a given query.

    Args:
        query: Natural language prompt to research.
        **kwargs: Optional configuration parameters:
            - embedding_model_name: Model for text embeddings (default: 'all-MiniLM-L6-v2')
            - llm_model: LLM model name (default: 'llama3' for Ollama, or HF model like 'meta-llama/Llama-2-7b-chat-hf')
            - llm_provider: 'ollama' or 'huggingface' (default: 'ollama')
            - hf_token: HuggingFace API token (required for gated models)
            - max_sources: Max search results to scrape (default: 5)
            - scraping_depth: Currently dictates one layer depth (default: 1)
            - style: Summarization style (default: 'concise')
            - research_mode: 'general', 'academic', 'documentation', 'github' (default: 'general')

    Returns:
        Dict containing the answer, references, and extracted information.
    """
    engine = WebResearchEngine(**kwargs)
    return engine.run(query)
