import subprocess
import sys
import json
import os
import re
import base64
import numpy as np
from io import BytesIO
from PIL import Image
import openai
from googlesearch import search
import requests
from bs4 import BeautifulSoup

# For the Scraper
import torch

api_key = os.getenv("OPENAI_API_KEY")

# Get the node list
script_dir = os.path.dirname(os.path.abspath(__file__))
json_file_path = os.path.join(script_dir, "../ComfyUI-Manager/custom-node-list.json")
json_file_path = os.path.normpath(json_file_path)

# Get the directory of the current script (Polymath.py)
script_directory = os.path.dirname(os.path.abspath(__file__))
custom_instructions_directory = os.path.join(script_directory, 'custom_instructions')

# Construct the full path to config.json
config_path = os.path.join(script_directory, 'config.json')

def install_and_import(package):
    try:
        __import__(package)
    except ImportError:
        print(f"Package {package} not found. Installing...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
    finally:
        globals()[package] = __import__(package)

install_and_import("googlesearch")
install_and_import("requests")
install_and_import("bs4")

script_directory = os.path.dirname(os.path.abspath(__file__))
config_path = os.path.join(script_directory, 'config.json')
ollama_url = "http://127.0.0.1:11434/api/tags"

def load_models():
    # Initialize dictionaries
    config_models = {}
    ollama_models = {}

    # Load models from config.json
    try:
        with open(config_path, 'r') as f:
            config_data = json.load(f)
            config_models = config_data.get('models', {})
            base_url = config_data.get('baseurl')
    except FileNotFoundError:
        print(f"Error: config.json not found at {config_path}")
    except json.JSONDecodeError:
        print(f"Error: Could not decode JSON from {config_path}")

    # Fetch models from Ollama API
    try:
        response = requests.get(ollama_url)
        response.raise_for_status()
        models = response.json().get('models', [])
        ollama_models = {model['name']: model['name'] for model in models}
    except requests.exceptions.RequestException as e:
        print(f"Error fetching Ollama models: {e}")

    # Merge dictionaries
    merged_models = {**config_models, **ollama_models}
     
    # Create a list of model names
    models_list = list(merged_models.keys())

    # Determine the default model
    default_model = models_list[0] if models_list else None
    print("\033[92mPolymath loaded these models: \033[0m", f"\033[38;5;99m{models_list}\033[0m")

    return base_url, merged_models, models_list, default_model

# Load models and set default
base_url, models_dict, models_list, default_model = load_models()

def load_custom_instructions():
    instructions = ["None"]
    if not os.path.exists(custom_instructions_directory):
        os.makedirs(custom_instructions_directory)
    for filename in os.listdir(custom_instructions_directory):
        if filename.endswith(".txt"):
            instructions.append(filename[:-4])
    return instructions

custom_instructions_list = load_custom_instructions()

def fetch_and_extract_content(url):
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        soup = BeautifulSoup(response.content, 'html.parser')
        paragraphs = soup.find_all('p')
        text_content = "\n".join([p.text for p in paragraphs])
        return text_content.strip()
    except requests.exceptions.RequestException as e:
        print(f"Error fetching URL {url}: {e}")
        return None
    except Exception as e:
        print(f"Error processing content from {url}: {e}")
        return None

class OllamaAPI:
    def __init__(self, base_url="http://localhost:11434"):
        self.base_url = base_url
        self.chat_history = []

    def generate_completion(self, model, prompt, stream=False, keep_alive=5, ollama_chat_mode=False, images=None):
        url = f"{self.base_url}/api/chat" if ollama_chat_mode else f"{self.base_url}/api/generate"

        if ollama_chat_mode:
            # Append the new user message to the chat history
            self.chat_history.append({"role": "user", "content": prompt})
            payload = {
                "model": model,
                "messages": self.chat_history,  # Include the entire chat history
                "stream": stream,
                "keep_alive": keep_alive
            }
        else:
            payload = {
                "model": model,
                "prompt": prompt,
                "stream": stream,
                "keep_alive": keep_alive
            }

        if images:
            payload["images"] = images

        response = requests.post(url, json=payload)
        response.raise_for_status()
        response_json = response.json()

        if ollama_chat_mode:
            # Append the model's response to the chat history
            self.chat_history.append({"role": "assistant", "content": response_json.get('message', {}).get('content', '')})

        return response_json

class Polymath:
    chat_history = []
    default_model = None
    ollama_api = OllamaAPI()  # Class-level instance of OllamaAPI

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": (
                    "STRING",
                    {
                        "default": "Enter your prompt here. Use {additional_text} as a placeholder if needed.",
                        "multiline": True
                    }
                ),
                "additional_text": (
                    "STRING",
                    {
                        "default": "",
                        "multiline": True
                    }
                ),
                "seed": ("INT", {"default": 0, "min": 0, "step": 1}),
            },
            "optional": {
                "model": (models_list, {"default": default_model}),
                "custom_instruction": (custom_instructions_list, {"default": "None"}),
                "enable_web_search": ("BOOLEAN", {"default": False}),
                "list_sources": ("BOOLEAN", {"default": False}),
                "num_search_results": ("INT", {"default": 5, "min": 1, "max": 10}),
                "keep_context": ("BOOLEAN", {"default": True}),
                "ollama_chat_mode": ("BOOLEAN", {"default": False}),
                "compress": ("BOOLEAN", {"default": False}),
                "compression_level": (["soft", "medium", "hard"],),
                "console_log": ("BOOLEAN", {"default": True}),
                "image": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "execute"
    CATEGORY = "Polymath"

    def tensor_to_pil(self, img_tensor):
        arr = img_tensor.cpu().numpy() if hasattr(img_tensor, "cpu") else np.array(img_tensor)
        arr = np.squeeze(arr)  # remove singleton dimensions
        if arr.ndim == 3 and arr.shape[0] == 3:  # if in (C, H, W), transpose to (H, W, C)
            arr = np.transpose(arr, (1, 2, 0))
        arr = np.clip(255.0 * arr, 0, 255).astype(np.uint8)
        return Image.fromarray(arr)

    def encode_image(self, image):
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")

    def execute(
        self, prompt, additional_text, seed, model=None, custom_instruction="None",
        enable_web_search=False, list_sources=False, num_search_results=3, keep_context=True, ollama_chat_mode=False, compress=False, compression_level="soft", console_log=True, image=None
    ):
        selected_model_value = models_dict.get(model)
        if not selected_model_value:
            print(f"Error: Model '{model}' not found in config.json. Using default.")
            selected_model_value = list(models_dict.values())[0] if models_dict else None
            if not selected_model_value:
                return ("Error: No valid model found.",)

        if console_log:
            print(
                f"\033[92mPolymath Chat Request: \033[0mModel='{model}', Prompt='{prompt}', Additional Text='{additional_text}'", 
                f"Seed='{seed}', Web Search Enabled={enable_web_search}, Keep Context={keep_context},"
                f"Ollama Chat Mode={ollama_chat_mode}, Custom Instruction='{custom_instruction}'"
            )

        # Process additional_text using the given prompt:
        # If the prompt contains the placeholder "{additional_text}", then replace it with the additional_text input.
        # Otherwise, if additional_text is provided, append it at the end.
        if "{additional_text}" in prompt:
            augmented_prompt = prompt.replace("{additional_text}", additional_text)
        elif additional_text.strip():
            augmented_prompt = prompt + "\n\n" + additional_text
        else:
            augmented_prompt = prompt

        if enable_web_search:
            urls_in_prompt = re.findall(r'(https?://\S+)', augmented_prompt)
            try:
                search_results_content = []
                searched_urls = set()
                for url in urls_in_prompt:
                    if url not in searched_urls:
                        print(f"Fetching content from URL in prompt: {url}")
                        content = fetch_and_extract_content(url)
                        if content:
                            search_results_content.append(
                                f"Source (from prompt): {url}\nContent:\n{content}\n---\n"
                            )
                        else:
                            search_results_content.append(
                                f"Source (from prompt): {url}\nCould not retrieve content.\n---\n"
                            )
                        searched_urls.add(url)
                for i, url in enumerate(search(augmented_prompt, num_results=num_search_results)):
                    if url not in searched_urls:
                        print(f"Fetching content from search result {i+1}: {url}")
                        content = fetch_and_extract_content(url)
                        if content:
                            search_results_content.append(
                                f"Search Result {i+1}:\nSource: {url}\nContent:\n{content}\n---\n"
                            )
                        else:
                            search_results_content.append(
                                f"Search Result {i+1}:\nSource: {url}\nCould not retrieve content.\n---\n"
                            )
                        searched_urls.add(url)
                if search_results_content:
                    augmented_prompt = (
                        f"Original query: {augmented_prompt}\n\nWeb search results and linked content:\n\n"
                        f"{''.join(search_results_content)}\n\nBased on this information, please provide a response."
                    )
                else:
                    augmented_prompt = (
                        f"Original query: {augmented_prompt}\n\nNo relevant web search results or linked content found. "
                        "Please proceed with the original query."
                    )
                    if console_log:
                        print("No relevant web search results or linked content found.")
            except Exception as e:
                print(f"Web Search Error: {e}")
                augmented_prompt = (
                    f"Original query: {augmented_prompt}\n\nAn error occurred during web search or fetching linked content. "
                    "Please proceed with the original query."
                )

        if compress:
            compression_chars = {
                "soft": 250,
                "medium": 150,
                "hard": 75,
            }
            char_limit = compression_chars[compression_level]
            augmented_prompt += (
                f" Compress the output to be concise while retaining key visual details. MAX OUTPUT SIZE no more than "
                f"{char_limit} characters."
            )
            if console_log:
                print(f"\033[92m{augmented_prompt}\033[0m")

        if list_sources:
            augmented_prompt += " Always list all fetched sources."

        if custom_instruction == "node finder":
            with open(json_file_path, "r", encoding="utf-8") as file:
                json_data = json.load(file)
                filtered_data = [
                    f"Title: {node['title']}\nReference: {node['reference']}\nDescription: {node['description']}\n"
                    for node in json_data.get("custom_nodes", [])
                ]
                data = "\n".join(filtered_data)
                print("Filtered JSON Data:")
                print(data)
                augmented_prompt += data

        b64 = None
        if image is not None:
            pil_image = self.tensor_to_pil(image)
            b64 = self.encode_image(pil_image)

        response = self.polymath_interaction(
            base_url, api_key, selected_model_value, augmented_prompt,
            console_log, keep_context, ollama_chat_mode, custom_instruction, b64=b64, 
        )
        return response

    def polymath_interaction(self, base_url, api_key, model_value, prompt, console_log, keep_context, ollama_chat_mode, custom_instruction, b64=None):
        
        if model_value.startswith(('gpt', 'o1', 'o3')):        
            client = openai.OpenAI(api_key=api_key, base_url=base_url)
            messages = []
            if custom_instruction != "None":
                instruction_path = os.path.join(custom_instructions_directory, f"{custom_instruction}.txt")
                try:
                    with open(instruction_path, 'r') as f:
                        system_instruction = f.read()
                        messages.append({"role": "system", "content": system_instruction})
                except FileNotFoundError:
                    print(f"Error: Custom instruction file '{instruction_path}' not found.")
                except Exception as e:
                    print(f"Error reading custom instruction file: {e}")

            if keep_context and self.chat_history:
                messages.extend(self.chat_history)
            if b64:
                image_data = {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}
                user_content = [
                    {"type": "text", "text": prompt},
                    image_data
                ]
                messages.append({"role": "user", "content": user_content})
            else:
                messages.append({"role": "user", "content": prompt})

            completion = client.chat.completions.create(
                model=model_value,
                messages=messages,
                stream=False
            )
            output_text = completion.choices[0].message.content

            if console_log:
                print("\033[92mPolymath Chat Response:\033[0m", output_text)
            self.chat_history.append({"role": "user", "content": prompt})
            self.chat_history.append({"role": "assistant", "content": output_text})

            return (output_text,)
        
        else:
            images = None

            if b64:
                images = [b64]

            response = self.ollama_api.generate_completion(
                ollama_chat_mode=ollama_chat_mode,
                model=model_value,
                prompt=prompt,
                stream=False,
                keep_alive=5,
                images=images
            )

            if ollama_chat_mode:
                output_text = response.get('message', {}).get('content', '')
            else:
                output_text = response.get('response', '')

            if console_log:
                print("\033[92mOllama API Response:\033[0m", output_text)
                
            return (output_text,)

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

class MediaScraper:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "urls": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": False}),
                "output_file_path": ("STRING", {"multiline": False, "default": ""}),
                "file_name": ("STRING", {"multiline": False, "default": ""}),
            },
            "optional": {
                "image_load_cap": ("INT", {"default": 0, "min": 0, "step": 1})
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("images", "filepath_texts")
    FUNCTION = "execute"
    CATEGORY = "Polymath"
    OUTPUT_IS_LIST = (True, True)

    def execute(self, urls, output_file_path, file_name, image_load_cap=0):
        url_list = [u.strip() for u in urls.split("\n") if u.strip()]
        downloaded_files = []

        # Download step: Run gallery-dl for each URL
        for url in url_list:
            if image_load_cap > 0 and len(downloaded_files) >= image_load_cap:
                break
            cmd = ["gallery-dl", url]
            if output_file_path:
                cmd.extend(["-d", output_file_path])
            if image_load_cap > 0:
                cap = image_load_cap - len(downloaded_files)
                cmd.extend(["--range", f"1-{cap}"])
            subprocess.run(cmd)

        # List all downloaded files from the output directory
        valid_extensions = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif')
        if output_file_path and os.path.isdir(output_file_path):
            for root, dirs, files in os.walk(output_file_path):
                for f in sorted(files):
                    if f.lower().endswith(valid_extensions):
                        full_path = os.path.join(root, f)
                        downloaded_files.append(full_path)
                        if image_load_cap > 0 and len(downloaded_files) >= image_load_cap:
                            break
                if image_load_cap > 0 and len(downloaded_files) >= image_load_cap:
                    break

        # Batch rename if file_name is provided
        if file_name and downloaded_files:
            base, ext = os.path.splitext(file_name)
            ext = ext if ext else ".jpg"
            for idx, old_path in enumerate(downloaded_files, start=1):
                new_name = f"{base}_{idx}{ext}"
                new_path = os.path.join(os.path.dirname(old_path), new_name)
                os.rename(old_path, new_path)
                downloaded_files[idx-1] = new_path

        valid_files = [f for f in downloaded_files if os.path.isfile(f) and f.lower().endswith(valid_extensions)]
        
        # Initialize empty lists for images and filepath_texts
        images = []
        filepath_texts = []

        for file_path in valid_files:
            try:
                img = Image.open(file_path)
                images.append(pil2tensor(img))
                filepath_texts.append(file_path)  # Add file path to the list
            except Exception:
                pass

        if not images:
            images.append(torch.ones(3, 100, 100))
            filepath_texts.append("")  # Add an empty string if no images are found
        
        return (images, filepath_texts)

NODE_CLASS_MAPPINGS = {
    "polymath_chat": Polymath,
    "polymath_scraper": MediaScraper
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "polymath_chat": "LLM Polymath Chat with Advanced Web and Link Search",
    "polymath_scraper": "LLM Polymath Scraper for various sites"
}