import mistune
from mistune.plugins.table import table
from jinja2 import Template
import re
import os
import hrequests
import markdown
from bs4 import BeautifulSoup
from lxml import etree
import markdown
import logging
from datetime import datetime
import psycopg2
from dotenv import load_dotenv
import ast
from fpdf import FPDF
import pandas as pd
import nltk
import requests
import json
from retry import retry
from concurrent.futures import ThreadPoolExecutor, as_completed
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from brave import Brave
from fuzzy_json import loads
from half_json.core import JSONFixer
from openai import OpenAI
from together import Together
from urllib.parse import urlparse
import trafilatura
import tiktoken

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Load environment variables
load_dotenv("keys.env")
TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HELICON_API_KEY = os.getenv("HELICON_API_KEY")
SUPABASE_USER = os.environ['SUPABASE_USER']
SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']
OPENROUTER_API_KEY = "sk-or-v1-" + os.environ['OPENROUTER_API_KEY']

# Define constants
LLM_DEFAULT_SMALL = "llama3-8b-8192"
LLM_DEFAULT_MEDIUM = "llama3-70b-8192"
LLM_FALLBACK_SMALL = "meta-llama/Llama-3-8b-chat-hf"
LLM_FALLBACK_MEDIUM = "meta-llama/Llama-3-70b-chat-hf"

SYS_PROMPT_DATA = """
You are an AI assistant tasked with extracting relevant information from scraped website data based on a given query. Your goal is to provide accurate and concise information that directly relates to the query, using only the data provided.
Guidelines for extraction:
1. Only use information present in the scraped data.
2. Focus on extracting facts, tables, and direct quotes that are relevant to the query.
3. If there is no relevant information in the scraped data, state that clearly.
4. Do not make assumptions or add information not present in the data.
5. If the query is ambiguous, interpret it in the most reasonable way based on the available data.
"""

SYS_PROMPT_DEFAULT = "You are an expert AI, complete the given task. Do not add any additional comments."
SYS_PROMPT_SEARCH = """You are a search query generator, create a concise Google search query, focusing only on the main topic and omitting additional redundant details, include year if necessary, 2024, Do not add any additional comments. OUTPUT ONLY THE SEARCH QUERY
#Additional instructions:
##Use the following search operator if necessary
OR #to cover multiple topics"""

# Initialize API clients
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")

together_client = OpenAI(
    api_key=TOGETHER_API_KEY, 
    base_url="https://together.hconeai.com/v1", 
    default_headers={"Helicone-Auth": f"Bearer {HELICON_API_KEY}"})

groq_client = OpenAI(
    api_key=GROQ_API_KEY, 
    base_url="https://groq.hconeai.com/openai/v1", 
    default_headers={"Helicone-Auth": f"Bearer {HELICON_API_KEY}"})

or_client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY)

def md_to_html(md_text):
    try:
        html_content = markdown.markdown(md_text, extensions=["extra"])
        return html_content.replace('\n', '')
    except Exception as e:
        logging.error(f"Error converting markdown to HTML: {e}")
        return md_text

def has_tables(html_string):
    try:
        soup = BeautifulSoup(html_string, 'lxml')
        if soup.find_all('table'):
            return True
        tree = etree.HTML(str(soup))
        return len(tree.xpath('//table')) > 0
    except Exception as e:
        logging.error(f"Error checking for tables: {e}")
        return False

def extract_data_from_tag(input_string, tag):
    try:
        pattern = f'<{tag}.*?>(.*?)</{tag}>'
        matches = re.findall(pattern, input_string, re.DOTALL)
        if matches:
            out = '\n'.join(match.strip() for match in matches)
            return out if len(out) <= 0.8 * len(input_string) else input_string
        return input_string
    except Exception as e:
        logging.error(f"Error extracting data from tag: {e}")
        return input_string

def insert_data(user_id, user_query, subtopic_query, response, html_report):
    try:
        with psycopg2.connect(
            dbname="postgres",
            user=SUPABASE_USER,
            password=SUPABASE_PASSWORD,
            host="aws-0-us-west-1.pooler.supabase.com",
            port="5432"
        ) as conn:
            with conn.cursor() as cur:
                insert_query = """
                INSERT INTO research_pro_chat_v2 (user_id, user_query, subtopic_query, response, html_report, created_at)
                VALUES (%s, %s, %s, %s, %s, %s);
                """
                cur.execute(insert_query, (user_id, user_query, subtopic_query, response, html_report, datetime.now()))
    except Exception as e:
        logging.error(f"Error inserting data into database: {e}")

def limit_tokens(input_string, token_limit=7500):
    try:
        return encoding.decode(encoding.encode(input_string)[:token_limit])
    except Exception as e:
        logging.error(f"Error limiting tokens: {e}")
        return input_string[:token_limit]  # Fallback to simple string slicing

def together_response(message, model=LLM_DEFAULT_SMALL, SysPrompt=SYS_PROMPT_DEFAULT, temperature=0.2, frequency_penalty=0.1, max_tokens=2000):
    messages = [{"role": "system", "content": SysPrompt}, {"role": "user", "content": message}]
    params = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "frequency_penalty": frequency_penalty,
        "max_tokens": max_tokens
    }
    try:
        response = groq_client.chat.completions.create(**params)
        return response.choices[0].message.content
    except Exception as e:
        logging.error(f"Error calling GROQ API: {e}")
        try:
            params["model"] = LLM_FALLBACK_SMALL if model == LLM_DEFAULT_SMALL else LLM_FALLBACK_MEDIUM 
            response = together_client.chat.completions.create(**params)
            return response.choices[0].message.content
        except Exception as e:
            logging.error(f"Error calling Together API: {e}")
            return "An error occurred while processing your request."

def openrouter_response(messages, model="meta-llama/llama-3-70b-instruct:nitro"):
    try:
        response = or_client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=4096,
        )
        return response.choices[0].message.content
    except Exception as e:
        logging.error(f"Error calling OpenRouter API: {e}")
        return None

def openrouter_response_stream(messages, model="meta-llama/llama-3-70b-instruct:nitro"):
    try:
        response = or_client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=4096,
            stream=True
        )
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                yield chunk.choices[0].delta.content
    except Exception as e:
        logging.error(f"Error streaming response from OpenRouter API: {e}")
        yield "An error occurred while streaming the response."

def json_from_text(text):
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        try:
            match = re.search(r'\{[\s\S]*\}', text)
            json_out = match.group(0) if match else text
            return loads(json_out)
        except Exception as e:
            logging.error(f"Error parsing JSON from text: {e}")
            return {}

def remove_stopwords(text):
    try:
        stop_words = set(stopwords.words('english'))
        words = word_tokenize(text)
        filtered_text = [word for word in words if word.lower() not in stop_words]
        return ' '.join(filtered_text)
    except Exception as e:
        logging.error(f"Error removing stopwords: {e}")
        return text

def rephrase_content(data_format, content, query):
    try:
        if data_format == "Structured data":
            return together_response(
                f"""return only the relevant information regarding the query: {{{query}}}. Output should be concise chunks of \
                paragraphs or tables or both, extracted from the following scraped context {{{limit_tokens(content,token_limit=2000)}}}""",
                SysPrompt=SYS_PROMPT_DATA,
                max_tokens=900,
            )
        elif data_format == "Quantitative data":
            return together_response(
                f"return only the numerical or quantitative data regarding the query: {{{query}}} structured into .md tables, using the scraped context:{{{limit_tokens(content,token_limit=2000)}}}",
                SysPrompt=SYS_PROMPT_DATA,
                max_tokens=500,
            )
        else:
            return together_response(
                f"return only the relevant information regarding the query: {{{query}}} using the scraped context:{{{limit_tokens(content,token_limit=2000)}}}",
                SysPrompt=SYS_PROMPT_DATA,
                max_tokens=500,
            )
    except Exception as e:
        logging.error(f"Error rephrasing content: {e}")
        return limit_tokens(content, token_limit=500)

def fetch_content(url):
    try:
        response = hrequests.get(url, timeout=5)
        if response.status_code == 200:
            return response.text
        else:
            logging.warning(f"Failed to fetch content from {url}. Status code: {response.status_code}")
    except Exception as e:
        logging.error(f"Error fetching page content for {url}: {e}")
    return None

def extract_main_content(html):
    try:
        extracted = trafilatura.extract(
            html,
            output_format="markdown",
            target_language="en",
            include_tables=True,
            include_images=False,
            include_links=False,
            deduplicate=True,
        )
        return trafilatura.utils.sanitize(extracted) if extracted else ""
    except Exception as e:
        logging.error(f"Error extracting main content: {e}")
        return ""

def process_content(data_format, url, query):
    try:
        html_content = fetch_content(url)
        if html_content:
            content = extract_main_content(html_content)
            if content:
                rephrased_content = rephrase_content(
                    data_format=data_format,
                    content=limit_tokens(remove_stopwords(content), token_limit=4000),
                    query=query,
                )
                return rephrased_content, url
    except Exception as e:
        logging.error(f"Error processing content for {url}: {e}")
    return "", url

def fetch_and_extract_content(data_format, urls, query):
    try:
        with ThreadPoolExecutor(max_workers=len(urls)) as executor:
            future_to_url = {
                executor.submit(process_content, data_format, url, query): url
                for url in urls
            }
            all_text_with_urls = [future.result() for future in as_completed(future_to_url)]
        return all_text_with_urls
    except Exception as e:
        logging.error(f"Error fetching and extracting content: {e}")
        return []

def search_brave(query, num_results=5):
    try:
        cleaned_query = query
        search_query = together_response(cleaned_query, model=LLM_DEFAULT_SMALL, SysPrompt=SYS_PROMPT_SEARCH, max_tokens=25).strip()
        cleaned_search_query = re.sub(r'[^\w\s]', '', search_query).strip()
        
        url = "https://api.search.brave.com/res/v1/web/search"
        headers = {
            "Accept": "application/json",
            "Accept-Encoding": "gzip",
            "X-Subscription-Token": BRAVE_API_KEY 
        }
        params = {"q": cleaned_search_query}
        
        response = requests.get(url, headers=headers, params=params)
        
        if response.status_code == 200:
            result = response.json()
            return [item["url"] for item in result["web"]["results"]][:num_results], cleaned_search_query, result
        else:
            logging.warning(f"Brave search API returned status code {response.status_code}")
            return [], cleaned_search_query, None
    except Exception as e:
        logging.error(f"Error in Brave search: {e}")
        return [], query, None

# Main execution
if __name__ == "__main__":
    logging.info("Script started")
    # Add your main execution logic here
    logging.info("Script completed")