"""
Data and Tool nodes for file I/O, web scraping, email, and databases.
"""

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

class WebScraperNode(BaseNode):
    """
    Node for scraping web pages.
    """
    
    def __init__(self, id: str, url: str = "", parser: str = "html.parser", **config):
        config.update({
            "url": url,
            "parser": parser
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.WEB_SCRAPER, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="url", type="str", required=False, description="URL to scrape (overrides config)")
        ]
        self.metadata.outputs = [
            NodeOutput(name="content", type="str", description="Main text content"),
            NodeOutput(name="html", type="str", description="Raw HTML"),
            NodeOutput(name="title", type="str", description="Page title"),
            NodeOutput(name="links", type="list", description="List of links")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        try:
            import httpx
            from bs4 import BeautifulSoup
        except ImportError:
            raise ImportError("Required packages not installed. Install with: pip install httpx beautifulsoup4")
        
        url = inputs.get("url") or self.metadata.config.get("url")
        if not url:
            raise ValueError("URL not provided")
        
        parser = self.metadata.config.get("parser", "html.parser")
        
        # Fetch page
        with httpx.Client(follow_redirects=True) as client:
            response = client.get(url, timeout=30.0)
            response.raise_for_status()
        
        soup = BeautifulSoup(response.text, parser)
        
        # Extract data
        title = soup.title.string if soup.title else ""
        
        # Get text content (remove scripts and styles)
        for script in soup(["script", "style"]):
            script.decompose()
        text = soup.get_text(separator="\n", strip=True)
        
        # Get links
        links = [a.get("href") for a in soup.find_all("a", href=True)]
        
        return {
            "content": text,
            "html": response.text,
            "title": title,
            "links": links
        }


class FileReadNode(BaseNode):
    """
    Node for reading files (TXT, CSV, JSON, Excel).
    """
    
    def __init__(self, id: str, filepath: str = "", file_type: str = "auto", **config):
        config.update({
            "filepath": filepath,
            "file_type": file_type
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.FILE_READ, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="filepath", type="str", required=False, description="Path to file")
        ]
        self.metadata.outputs = [
            NodeOutput(name="data", type="any", description="File content")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        filepath = inputs.get("filepath") or self.metadata.config.get("filepath")
        if not filepath:
            raise ValueError("Filepath not provided")
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")
        
        file_type = self.metadata.config.get("file_type", "auto").lower()
        if file_type == "auto":
            ext = os.path.splitext(filepath)[1].lower()
            if ext == ".csv":
                file_type = "csv"
            elif ext == ".json":
                file_type = "json"
            elif ext in [".xls", ".xlsx"]:
                file_type = "excel"
            else:
                file_type = "text"
        
        if file_type == "text":
            with open(filepath, "r", encoding="utf-8") as f:
                data = f.read()
        elif file_type == "json":
            with open(filepath, "r", encoding="utf-8") as f:
                data = json.load(f)
        elif file_type in ["csv", "excel"]:
            try:
                import pandas as pd
            except ImportError:
                raise ImportError("pandas not installed. Install with: pip install pandas")
            
            if file_type == "csv":
                df = pd.read_csv(filepath)
            else:
                df = pd.read_excel(filepath)
            
            # Convert to dict records for easier usage
            data = df.to_dict(orient="records")
        else:
            raise ValueError(f"Unsupported file type: {file_type}")
        
        return {"data": data}


class FileWriteNode(BaseNode):
    """
    Node for writing files.
    """
    
    def __init__(self, id: str, filepath: str = "", mode: str = "w", **config):
        config.update({
            "filepath": filepath,
            "mode": mode
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.FILE_WRITE, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="filepath", type="str", required=False, description="Path to file"),
            NodeInput(name="data", type="any", required=True, description="Data to write")
        ]
        self.metadata.outputs = [
            NodeOutput(name="success", type="bool", description="Operation success"),
            NodeOutput(name="filepath", type="str", description="Path written to")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        filepath = inputs.get("filepath") or self.metadata.config.get("filepath")
        if not filepath:
            raise ValueError("Filepath not provided")
        
        data = inputs.get("data")
        mode = self.metadata.config.get("mode", "w")
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
        
        if isinstance(data, (dict, list)):
            # Write as JSON if data is structured and file extension is .json
            if filepath.endswith(".json"):
                with open(filepath, "w", encoding="utf-8") as f:
                    json.dump(data, f, indent=2)
            elif filepath.endswith(".csv") and isinstance(data, list):
                try:
                    import pandas as pd
                    df = pd.DataFrame(data)
                    df.to_csv(filepath, index=False)
                except ImportError:
                    # Fallback to string
                    with open(filepath, mode, encoding="utf-8") as f:
                        f.write(str(data))
            else:
                with open(filepath, mode, encoding="utf-8") as f:
                    f.write(str(data))
        else:
            with open(filepath, mode, encoding="utf-8") as f:
                f.write(str(data))
        
        return {"success": True, "filepath": filepath}


class EmailNode(BaseNode):
    """
    Node for sending emails via SMTP.
    """
    
    def __init__(
        self, 
        id: str, 
        smtp_server: str = "", 
        port: int = 587, 
        username: str = "", 
        password: str = "",
        **config
    ):
        config.update({
            "smtp_server": smtp_server,
            "port": port,
            "username": username,
            "password": password
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.EMAIL, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="to", type="str", required=True, description="Recipient email"),
            NodeInput(name="subject", type="str", required=True, description="Email subject"),
            NodeInput(name="body", type="str", required=True, description="Email body"),
            NodeInput(name="from", type="str", required=False, description="Sender email (optional)")
        ]
        self.metadata.outputs = [
            NodeOutput(name="success", type="bool", description="Operation success")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        import smtplib
        from email.mime.text import MIMEText
        from email.mime.multipart import MIMEMultipart
        
        smtp_server = self.metadata.config.get("smtp_server")
        port = self.metadata.config.get("port", 587)
        username = self.metadata.config.get("username")
        password = self.metadata.config.get("password")
        
        if not smtp_server or not username or not password:
            # Try env vars
            smtp_server = smtp_server or os.environ.get("SMTP_SERVER")
            port = port or int(os.environ.get("SMTP_PORT", 587))
            username = username or os.environ.get("SMTP_USERNAME")
            password = password or os.environ.get("SMTP_PASSWORD")
            
            if not smtp_server or not username or not password:
                raise ValueError("SMTP configuration missing")
        
        to_email = inputs.get("to")
        subject = inputs.get("subject")
        body = inputs.get("body")
        from_email = inputs.get("from") or username
        
        msg = MIMEMultipart()
        msg["From"] = from_email
        msg["To"] = to_email
        msg["Subject"] = subject
        msg.attach(MIMEText(body, "plain"))
        
        with smtplib.SMTP(smtp_server, port) as server:
            server.starttls()
            server.login(username, password)
            server.send_message(msg)
        
        return {"success": True}


class DatabaseNode(BaseNode):
    """
    Node for executing SQL queries.
    """
    
    def __init__(self, id: str, connection_string: str = "", **config):
        config.update({
            "connection_string": connection_string
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.DATABASE, **config)
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="query", type="str", required=True, description="SQL query"),
            NodeInput(name="params", type="dict", required=False, description="Query parameters")
        ]
        self.metadata.outputs = [
            NodeOutput(name="results", type="list", description="Query results"),
            NodeOutput(name="row_count", type="int", description="Affected rows")
        ]
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        try:
            from sqlalchemy import create_engine, text
        except ImportError:
            raise ImportError("sqlalchemy not installed. Install with: pip install sqlalchemy")
        
        conn_str = self.metadata.config.get("connection_string")
        if not conn_str:
            conn_str = os.environ.get("DB_CONNECTION_STRING")
            if not conn_str:
                raise ValueError("Connection string not provided")
        
        query = inputs.get("query")
        params = inputs.get("params", {})
        
        engine = create_engine(conn_str)
        
        with engine.connect() as conn:
            result = conn.execute(text(query), params)
            
            if result.returns_rows:
                rows = [dict(row._mapping) for row in result]
                return {"results": rows, "row_count": len(rows)}
            else:
                conn.commit()
                return {"results": [], "row_count": result.rowcount}


class JsonLoaderNode(FileReadNode):
    """
    Specialized node for loading JSON files and preparing them for vector search.
    """
    
    def __init__(self, id: str, filepath: str = "", text_key: str = None, **config):
        config["text_key"] = text_key
        super().__init__(id=id, filepath=filepath, file_type="json", **config)

    def _define_inputs_outputs(self):
        super()._define_inputs_outputs()
        self.metadata.outputs.extend([
            NodeOutput(name="documents", type="list", description="Extracted text documents"),
            NodeOutput(name="metadatas", type="list", description="Extracted metadata")
        ])

    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        result = super().execute(inputs, context)
        data = result["data"]
        text_key = self.metadata.config.get("text_key")
        
        documents = []
        metadatas = []
        
        if isinstance(data, list):
            for item in data:
                if isinstance(item, dict) and text_key:
                    # Extract text using key
                    doc = item.get(text_key, "")
                    documents.append(str(doc))
                    # Use rest as metadata
                    meta = item.copy()
                    meta.pop(text_key, None)
                    # Ensure metadata values are simple types (str, int, float, bool) for Chroma
                    clean_meta = {}
                    for k, v in meta.items():
                        if isinstance(v, (str, int, float, bool)):
                            clean_meta[k] = v
                        else:
                            clean_meta[k] = str(v)
                    metadatas.append(clean_meta)
                elif isinstance(item, str):
                    documents.append(item)
                    metadatas.append({})
                else:
                    documents.append(str(item))
                    metadatas.append({})
        elif isinstance(data, dict):
            # Single document
            if text_key and text_key in data:
                documents.append(str(data[text_key]))
                meta = data.copy()
                meta.pop(text_key, None)
                clean_meta = {}
                for k, v in meta.items():
                    if isinstance(v, (str, int, float, bool)):
                        clean_meta[k] = v
                    else:
                        clean_meta[k] = str(v)
                metadatas.append(clean_meta)
            else:
                documents.append(json.dumps(data))
                metadatas.append({})
        else:
            documents.append(str(data))
            metadatas.append({})
            
        return {
            "data": data,
            "documents": documents,
            "metadatas": metadatas
        }

