import streamlit as st
import time
from langchain.agents import initialize_agent, Tool
from langchain.agents import AgentType
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import PromptTemplate
import requests

class LLM:
    def __init__(self, tools, model, search):
        if not isinstance(tools, list) or not all(isinstance(t, Tool) for t in tools):
            raise ValueError("Tools must be a list of Tool objects")
            
        web_search_available = any(
            tool.name == "Web Search" for tool in tools
        )
        if not web_search_available:
            raise ValueError("Web Search tool must be included in tools list")
            
        self.tools = tools
        self.model = model
        self.search = search
        
        self.agent = initialize_agent(
            self.tools,
            self.model,
            agent_type=AgentType.SELF_ASK_WITH_SEARCH,
            verbose=True,
            max_iterations=5,
            handle_parsing_errors=True,
            early_stopping_method="generate",
            memory=ConversationBufferWindowMemory(k=1)
        )

    def perform_web_search(self, query, max_retries=5, delay=1, timeout=8):
        retries = 0
        while retries < max_retries:
            try:
                search_results = self.search.run(query, timeout=timeout)
                if search_results:
                    return search_results[:1500]
            except requests.exceptions.Timeout:
                retries += 1
                st.warning(f"Web search timed out. Retrying ({retries}/{max_retries})...")
                time.sleep(delay)
            except Exception as e:
                retries += 1
                st.warning(f"Web search failed. Retrying ({retries}/{max_retries})... Error: {e}")
                time.sleep(delay)
        return "NaN"

    def get_llm_response(self, entity, query_type, web_results):
        prompt = PromptTemplate(
            template="""
You are a highly skilled information extractor. Your job is to extract the most relevant {query_type} from the following Web Search Results. 
Provide the exact value requested and return only that value—no explanations, context, or irrelevant information.

Entity: {entity}
Information to Extract: {query_type}

Web Search Results:
{web_results}

If you cannot find relevant information, return "NaN". Do not return anything else.

Your extracted response:
""",
            input_variables=["entity", "query_type", "web_results"]
        )

        try:
            response = self.agent.invoke({
                "input": prompt.format(
                    query_type=query_type,
                    entity=entity,
                    web_results=web_results,
                )
            })

            raw_response = response.get("output", "").strip()
            if raw_response:
                return raw_response
            else:
                return "NaN"
        
        except Exception as e:
            st.error(f"Error processing response: {str(e)}")
            return "NaN"

    def refine_answer_with_searches(self, entity, query_type, max_retries=2):
        search_query = f"{entity} current {query_type}"
        search_results = self.perform_web_search(search_query)
        extracted_answer = self.get_llm_response(entity, query_type, search_results)

        if extracted_answer == "NaN" and max_retries > 0:
            alternative_query = f"{entity} {query_type} detailed information"
            search_results = self.perform_web_search(alternative_query)
            extracted_answer = self.get_llm_response(entity, query_type, search_results)

        return extracted_answer, search_results