"""
Advanced AI nodes for Vector Search, RAG, and specialized NLP tasks.
"""

from typing import Any, Dict, List, Optional
from ybagent.node import BaseNode, NodeType, NodeInput, NodeOutput
import json
import os

class VectorSearchNode(BaseNode):
    """
    Node for Vector Search (RAG) using ChromaDB.
    """
    
    def __init__(self, id: str, collection_name: str = "ybagent_docs", **config):
        config.update({
            "collection_name": collection_name
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.VECTOR_SEARCH, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="query", type="str", required=True, description="Search query"),
            NodeInput(name="documents", type="list", required=False, description="Documents to add (list of strings)"),
            NodeInput(name="metadatas", type="list", required=False, description="Metadata for documents"),
            NodeInput(name="ids", type="list", required=False, description="IDs for documents")
        ]
        self.metadata.outputs = [
            NodeOutput(name="results", type="list", description="Top matching documents"),
            NodeOutput(name="distances", type="list", description="Similarity scores")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        try:
            import chromadb
            from chromadb.utils import embedding_functions
        except ImportError:
            raise ImportError("chromadb not installed. Install with: pip install chromadb sentence-transformers")
        
        meta_cfg = self.metadata.config.get("metadata")
        persist_dir = self.metadata.config.get("persist_directory")
        if isinstance(meta_cfg, dict) and not persist_dir:
            persist_dir = meta_cfg.get("persist_directory")
        persist_path = persist_dir or os.path.join(os.getcwd(), "chroma_db")
        client = chromadb.PersistentClient(path=persist_path)
        
        collection_name = self.metadata.config.get("collection_name", "ybagent_docs")
        
        embed_name = self.metadata.config.get("embedding_model_name")
        ef = None
        if embed_name:
            try:
                ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=embed_name)
            except Exception:
                ef = embedding_functions.DefaultEmbeddingFunction()
        else:
            ef = embedding_functions.DefaultEmbeddingFunction()
        
        collection = client.get_or_create_collection(name=collection_name, embedding_function=ef)
        
        documents = inputs.get("documents")
        if documents:
            if isinstance(documents, dict):
                if "documents" in documents and isinstance(documents["documents"], list):
                    documents = documents["documents"]
                else:
                    def _collect_strings(obj):
                        res = []
                        if isinstance(obj, str):
                            res.append(obj)
                        elif isinstance(obj, list):
                            for v in obj:
                                res.extend(_collect_strings(v))
                        elif isinstance(obj, dict):
                            for v in obj.values():
                                res.extend(_collect_strings(v))
                        return res
                    documents = _collect_strings(documents)
            elif not isinstance(documents, list):
                documents = [str(documents)]
            else:
                documents = [d if isinstance(d, str) else str(d) for d in documents]

            if documents:
                ids = inputs.get("ids")
                if not ids:
                    import uuid
                    ids = [str(uuid.uuid4()) for _ in range(len(documents))]
                metadatas = inputs.get("metadatas")
                collection.add(documents=documents, metadatas=metadatas, ids=ids)
        
        # Query
        query = inputs.get("query")
        if not query:
            return {"results": [], "distances": []}
        
        if isinstance(query, dict):
            def _first_string(val):
                if isinstance(val, str):
                    return val
                if isinstance(val, dict):
                    for k in ("query", "text", "input", "prompt", "value"):
                        if k in val and isinstance(val[k], str):
                            return val[k]
                    for v in val.values():
                        s = _first_string(v)
                        if s:
                            return s
                    return None
                if isinstance(val, list):
                    for v in val:
                        s = _first_string(v)
                        if s:
                            return s
                    return None
                return None
            s = _first_string(query)
            query_text = s if s is not None else str(query)
        else:
            query_text = str(query)
        
        n_results = self.metadata.config.get("n_results", 3)
        
        results = collection.query(
            query_texts=[query_text],
            n_results=n_results
        )
        
        # Flatten results
        flat_results = results["documents"][0] if results["documents"] else []
        flat_distances = results["distances"][0] if results["distances"] else []
        flat_metadatas = results["metadatas"][0] if results["metadatas"] else []
        
        return {
            "results": flat_results,
            "distances": flat_distances
        }


class BaseLLMTaskNode(BaseNode):
    """
    Base class for specialized LLM tasks (Summarization, Translation, etc.).
    Defaults to Ollama for execution.
    """
    
    def _execute_llm(self, prompt: str, system_prompt: str = "") -> str:
        provider = self.metadata.config.get("provider", "ollama")
        model = self.metadata.config.get("model", "llama2")
        
        if provider == "ollama":
            return self._execute_ollama(prompt, system_prompt, model)
        elif provider == "openai":
            return self._execute_openai(prompt, system_prompt, model)
        else:
            raise ValueError(f"Unsupported provider: {provider}")

    def _execute_ollama(self, prompt: str, system_prompt: str, model: str) -> str:
        import httpx
        
        url = self.metadata.config.get("base_url", "http://localhost:11434")
        
        payload = {
            "model": model,
            "prompt": prompt,
            "system": system_prompt,
            "stream": False,
            "options": {
                "temperature": self.metadata.config.get("temperature", 0.7)
            }
        }
        
        try:
            response = httpx.post(f"{url}/api/generate", json=payload, timeout=60.0)
            response.raise_for_status()
            return response.json().get("response", "")
        except Exception as e:
            raise RuntimeError(f"Ollama execution failed: {e}")

    def _execute_openai(self, prompt: str, system_prompt: str, model: str) -> str:
        from openai import OpenAI
        
        api_key = self.metadata.config.get("api_key") or os.environ.get("OPENAI_API_KEY")
        client = OpenAI(api_key=api_key)
        
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})
        
        try:
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=self.metadata.config.get("temperature", 0.7)
            )
            return response.choices[0].message.content
        except Exception as e:
            raise RuntimeError(f"OpenAI execution failed: {e}")


class SummarizationNode(BaseLLMTaskNode):
    """
    Node for summarizing text.
    """
    
    def __init__(self, id: str, **config):
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.SUMMARIZATION, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="text", type="str", required=True, description="Text to summarize")
        ]
        self.metadata.outputs = [
            NodeOutput(name="summary", type="str", description="Summarized text")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        text = inputs.get("text", "")
        if not text:
            raise ValueError("No text provided for summarization")
        
        system_prompt = "You are a helpful assistant that summarizes text concisely."
        prompt = f"Please summarize the following text:\n\n{text}"
        
        summary = self._execute_llm(prompt, system_prompt)
        return {"summary": summary}


class TranslationNode(BaseLLMTaskNode):
    """
    Node for translating text.
    """
    
    def __init__(self, id: str, source_lang: str = "auto", target_lang: str = "English", **config):
        config.update({
            "source_lang": source_lang,
            "target_lang": target_lang
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.TRANSLATION, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="text", type="str", required=True, description="Text to translate"),
            NodeInput(name="llm_response", type="str", required=False, description="External LLM translation output")
        ]
        self.metadata.outputs = [
            NodeOutput(name="translated_text", type="str", description="Translated text"),
            NodeOutput(name="prompt", type="str", description="Prompt for external LLM"),
            NodeOutput(name="system_message", type="str", description="System message for external LLM")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        text = inputs.get("text", "")
        source = self.metadata.config.get("source_lang", "auto")
        target = self.metadata.config.get("target_lang", "English")
        system_prompt = f"You are a professional translator. Translate from {source} to {target}."
        prompt = f"Translate:\n{text}"
        external = inputs.get("llm_response")
        if external is not None:
            return {"translated_text": external, "prompt": prompt, "system_message": system_prompt}
        provider = self.metadata.config.get("provider")
        if not provider:
            return {"translated_text": "", "prompt": prompt, "system_message": system_prompt}
        result = self._execute_llm(prompt, system_prompt)
        return {"translated_text": result, "prompt": prompt, "system_message": system_prompt}


class SentimentAnalysisNode(BaseLLMTaskNode):
    """
    Node for sentiment analysis.
    """
    
    def __init__(self, id: str, **config):
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.SENTIMENT, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="text", type="str", required=True, description="Text to analyze")
        ]
        self.metadata.outputs = [
            NodeOutput(name="sentiment", type="str", description="POSITIVE, NEGATIVE, or NEUTRAL"),
            NodeOutput(name="explanation", type="str", description="Reasoning")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        text = inputs.get("text", "")
        
        system_prompt = "You are a sentiment analyzer. Respond with JSON only: {\"sentiment\": \"POSITIVE/NEGATIVE/NEUTRAL\", \"explanation\": \"...\"}"
        prompt = f"Analyze sentiment:\n{text}"
        
        result = self._execute_llm(prompt, system_prompt)
        
        # Try to parse JSON
        try:
            # Clean up markdown code blocks if present
            if "```json" in result:
                result = result.split("```json")[1].split("```")[0]
            elif "```" in result:
                result = result.split("```")[1].split("```")[0]
                
            data = json.loads(result)
            return {
                "sentiment": data.get("sentiment", "UNKNOWN"),
                "explanation": data.get("explanation", "")
            }
        except:
            # Fallback
            return {"sentiment": "UNKNOWN", "explanation": result}


class ClassificationNode(BaseLLMTaskNode):
    """
    Node for text classification.
    """
    
    def __init__(self, id: str, categories: List[str] = None, **config):
        config.update({
            "categories": categories or []
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.CLASSIFICATION, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="text", type="str", required=True, description="Text to classify")
        ]
        self.metadata.outputs = [
            NodeOutput(name="category", type="str", description="Predicted category")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        text = inputs.get("text", "")
        categories = self.metadata.config.get("categories", [])
        
        if not categories:
            raise ValueError("No categories provided for classification")
            
        cats_str = ", ".join(categories)
        system_prompt = f"Classify the text into one of these categories: {cats_str}. Respond ONLY with the category name."
        prompt = f"Text:\n{text}"
        
        category = self._execute_llm(prompt, system_prompt).strip()
        return {"category": category}


class NERNode(BaseLLMTaskNode):
    """
    Node for Named Entity Recognition.
    """
    
    def __init__(self, id: str, **config):
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.NER, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="text", type="str", required=True, description="Text to analyze"),
            NodeInput(name="llm_response", type="str", required=False, description="External LLM NER output")
        ]
        self.metadata.outputs = [
            NodeOutput(name="entities", type="dict", description="Extracted entities (people, places, etc.)"),
            NodeOutput(name="prompt", type="str", description="Prompt for external LLM"),
            NodeOutput(name="system_message", type="str", description="System message for external LLM")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        text = inputs.get("text", "")
        system_prompt = self.metadata.config.get("system_prompt")
        if not system_prompt:
            raise ValueError("system_prompt must be provided in node configuration")
        prompt = f"Text:\n{text}"
        original_query = text
        if isinstance(original_query, dict):
            # Try to extract meaningful text from dict
            if "text" in original_query:
                original_query = original_query["text"]
            elif "input" in original_query:
                original_query = original_query["input"]
            elif "content" in original_query:
                original_query = original_query["content"]
            else:
                # Fallback to string representation but keep it clean
                original_query = str(original_query)
        
        # Ensure it looks like "Text: ..." for display
        if not str(original_query).strip().startswith("Text:"):
            display_query = f'Text: "{original_query}"'
        else:
            display_query = original_query
        external = inputs.get("llm_response")
        if external is not None:
            try:
                from ybagent.utils.display import DisplayManager
                state = context.get("state") if isinstance(context, dict) else None
                suppress_llm = state.get_variable("suppress_llm_response", False) if state else False
                if not suppress_llm:
                    DisplayManager.llm_response(str(external), title="NER Response", query=display_query)
            except ImportError:
                pass
            try:
                r = external
                if "```json" in r:
                    r = r.split("```json")[1].split("```")[0]
                elif "```" in r:
                    r = r.split("```")[1].split("```")[0]
                entities = json.loads(r)
                return {"entities": entities, "prompt": prompt, "system_message": system_prompt}
            except:
                return {"entities": {"raw": external}, "prompt": prompt, "system_message": system_prompt}
        provider = self.metadata.config.get("provider")
        if not provider:
            return {"entities": {}, "prompt": prompt, "system_message": system_prompt}
        result = self._execute_llm(prompt, system_prompt)
        try:
            from ybagent.utils.display import DisplayManager
            state = context.get("state") if isinstance(context, dict) else None
            suppress_llm = state.get_variable("suppress_llm_response", False) if state else False
            if not suppress_llm:
                DisplayManager.llm_response(str(result), title="NER Response", query=display_query)
        except ImportError:
            pass
        
        try:
            if "```json" in result:
                result = result.split("```json")[1].split("```")[0]
            elif "```" in result:
                result = result.split("```")[1].split("```")[0]
            entities = json.loads(result)
            return {"entities": entities, "prompt": prompt, "system_message": system_prompt}
        except:
            return {"entities": {"raw": result}, "prompt": prompt, "system_message": system_prompt}


class VisionNode(BaseLLMTaskNode):
    """
    Node for Vision tasks (Image/PDF analysis) using Multimodal LLMs.
    """
    
    def __init__(self, id: str, **config):
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.VISION, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="file_path", type="str", required=True, description="Path to image or PDF"),
            NodeInput(name="query", type="str", required=True, description="Question about the image/PDF"),
            NodeInput(name="system_prompt", type="str", required=False, description="Optional system instruction")
        ]
        self.metadata.outputs = [
            NodeOutput(name="response", type="str", description="LLM analysis response"),
            NodeOutput(name="images_processed", type="int", description="Number of images processed")
        ]
    
    def _encode_image(self, image_path: str) -> str:
        import base64
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')

    def _convert_pdf_to_images(self, pdf_path: str) -> List[str]:
        """Convert PDF pages to temporary images and return their paths."""
        try:
            # Try pdf2image (requires poppler)
            from pdf2image import convert_from_path
            import tempfile
            
            images = convert_from_path(pdf_path)
            image_paths = []
            
            for i, image in enumerate(images):
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
                    image.save(tmp.name, 'JPEG')
                    image_paths.append(tmp.name)
            return image_paths
            
        except ImportError:
            # Try pymupdf (fitz)
            try:
                import fitz  # PyMuPDF
                import tempfile
                
                doc = fitz.open(pdf_path)
                image_paths = []
                
                for page in doc:
                    pix = page.get_pixmap()
                    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                        pix.save(tmp.name)
                        image_paths.append(tmp.name)
                return image_paths
                
            except ImportError:
                raise ImportError(
                    "PDF processing requires 'pdf2image' (with poppler) or 'pymupdf'. "
                    "Please install one: pip install pdf2image OR pip install pymupdf"
                )

    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        file_path = inputs.get("file_path", "")
        query = inputs.get("query", "")
        system_prompt = inputs.get("system_prompt") or self.metadata.config.get("system_prompt", "You are a helpful vision assistant.")
        
        if not file_path or not os.path.exists(file_path):
            raise ValueError(f"File not found: {file_path}")
            
        # Determine file type
        ext = os.path.splitext(file_path)[1].lower()
        image_paths = []
        temp_files = []
        
        if ext in ['.pdf']:
            print(f"📄 Processing PDF: {file_path}")
            image_paths = self._convert_pdf_to_images(file_path)
            temp_files = image_paths # Track for cleanup
        elif ext in ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.webp']:
            image_paths = [file_path]
        else:
            raise ValueError(f"Unsupported file type: {ext}")
            
        # Encode images
        encoded_images = []
        for img_path in image_paths:
            encoded_images.append(self._encode_image(img_path))
            
        # Execute LLM
        # Allow overriding config from inputs
        provider = inputs.get("provider") or self.metadata.config.get("provider", "ollama")
        model = inputs.get("model") or self.metadata.config.get("model", "llava")
        api_key = inputs.get("api_key") or self.metadata.config.get("api_key") or os.environ.get("OPENAI_API_KEY")
        
        # Extract generation parameters
        temperature = float(inputs.get("temperature") or self.metadata.config.get("temperature", 0.7))
        max_tokens = int(inputs.get("max_tokens") or self.metadata.config.get("max_tokens", 1000))
        
        response_text = ""
        
        try:
            if provider == "ollama":
                import httpx
                url = self.metadata.config.get("base_url", "http://localhost:11434")
                
                # Ollama supports multiple images in one request
                payload = {
                    "model": model,
                    "prompt": query,
                    "system": system_prompt,
                    "images": encoded_images,
                    "stream": False,
                    "options": {
                        "temperature": temperature,
                        "num_predict": max_tokens
                    }
                }
                
                try:
                    response = httpx.post(f"{url}/api/generate", json=payload, timeout=120.0)
                    if response.status_code != 200:
                        raise RuntimeError(f"Ollama Error ({response.status_code}): {response.text}")
                    response.raise_for_status()
                    response_text = response.json().get("response", "")
                except Exception as e:
                    raise RuntimeError(f"Ollama vision execution failed: {e}")
                    
            elif provider == "openai":
                # OpenAI GPT-4 Vision
                from openai import OpenAI
                # api_key is already resolved above
                client = OpenAI(api_key=api_key)
                
                messages = [
                    {"role": "system", "content": system_prompt},
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": query},
                        ]
                    }
                ]
                
                # Add images
                for b64_img in encoded_images:
                    messages[1]["content"].append({
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{b64_img}"
                        }
                    })
                
                try:
                    response = client.chat.completions.create(
                        model=model,
                        messages=messages,
                        max_tokens=max_tokens,
                        temperature=temperature
                    )
                    response_text = response.choices[0].message.content
                except Exception as e:
                    raise RuntimeError(f"OpenAI vision execution failed: {e}")
            else:
                raise ValueError(f"Provider {provider} not supported for VisionNode")
                
            # Display result
            try:
                from ybagent.utils.display import DisplayManager
                state = context.get("state") if isinstance(context, dict) else None
                suppress_llm = state.get_variable("suppress_llm_response", False) if state else False
                
                if not suppress_llm:
                    # Clean response from code blocks
                    clean_response = response_text
                    if "```" in clean_response:
                        # Remove code block markers
                        import re
                        clean_response = re.sub(r'```[\w]*\n?', '', clean_response)
                        clean_response = clean_response.strip()
                    
                    # Format with Image and Query on separate lines
                    display_query = f'Image:\n"{os.path.basename(file_path)}"\n\nQuery:\n{query}'
                    DisplayManager.llm_response(clean_response, title="Vision Response", query=display_query)
            except ImportError:
                pass
                
        finally:
            # Cleanup temp files
            for tmp in temp_files:
                try:
                    os.remove(tmp)
                except:
                    pass
                    
        return {
            "response": response_text,
            "images_processed": len(encoded_images)
        }
