import requests
from typing import Any, Dict, List, Optional
from langchain.llms.base import LLM
from langchain.agents import initialize_agent, AgentType, Tool
from pydantic import Field
import os
from dotenv import load_dotenv
from datetime import datetime
import wikipedia
from asteval import Interpreter  # For a safer calculator
import logging
from .tools import tools_registry

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

load_dotenv()

LM_STUDIO_API_URL = os.getenv("LM_STUDIO_API_URL", "http://192.168.0.104:1234/v1/chat/completions")
MODEL_NAME = os.getenv("LM_STUDIO_MODEL", "lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf")
CONTENT_TYPE = "application/json"

class LMStudioLLM(LLM):
    """
    Custom LangChain LLM to interface with LM Studio API.
    """
    api_url: str = Field(default=LM_STUDIO_API_URL, description="The API endpoint for LM Studio.")
    model: str = Field(default=MODEL_NAME, description="The model path/name.")
    temperature: float = Field(default=0.7, description="Sampling temperature.")
    max_tokens: Optional[int] = Field(default=4096, description="Maximum number of tokens to generate.")
    streaming: bool = Field(default=False, alias="stream", description="Whether to use streaming responses.")

    class Config:
        populate_by_name = True

    @property
    def _llm_type(self) -> str:
        return "lmstudio"

    @property
    def identifying_params(self) -> Dict[str, Any]:
        return {
            "api_url": self.api_url,
            "model": self.model,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "stream": self.streaming,
        }

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        """
        Generate a response from the LM Studio model.

        Args:
            prompt (str): The input prompt.
            stop (Optional[List[str]]): Stop sequences.

        Returns:
            str: The generated response.
        """
        headers = {
            "Content-Type": CONTENT_TYPE,
        }

        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt},
            ],
            "temperature": self.temperature,
            "max_tokens": self.max_tokens if self.max_tokens is not None else -1,
            "stream": self.streaming,  # Uses alias 'stream'
        }

        logger.info(f"Payload: {payload}")

        try:
            response = requests.post(
                self.api_url,
                headers=headers,
                json=payload,
                timeout=60,
                stream=self.streaming
            )
            response.raise_for_status()
            logger.info(f"Response content: {response.text}")
        except requests.RequestException as e:
            logger.error(f"Failed to connect to LM Studio API: {e}")
            raise RuntimeError(f"Failed to connect to LM Studio API: {e}")

        if self.streaming:
            return self._handle_stream(response)
        else:
            try:
                response_json = response.json()
                choices = response_json.get("choices", [])
                if not choices:
                    raise ValueError("No choices found in the response.")

                # Extract the first response's content
                content = choices[0].get("message", {}).get("content", "")
                return content.strip()
            except (ValueError, KeyError) as e:
                logger.error(f"Invalid response format: {e}")
                raise RuntimeError(f"Invalid response format: {e}")

    def _handle_stream(self, response: requests.Response) -> str:
        """
        Process streaming responses from the LM Studio API.

        Args:
            response (requests.Response): The streaming response object.

        Returns:
            str: The concatenated content from the stream.
        """
        content = ""
        try:
            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith("data: "):
                        data = decoded_line[6:]
                        if data == "[DONE]":
                            break
                        try:
                            json_data = requests.utils.json.loads(data)
                            choices = json_data.get("choices", [])
                            for chunk in choices:
                                delta = chunk.get("delta", {})
                                piece = delta.get("content", "")
                                content += piece
                        except requests.utils.json.JSONDecodeError:
                            continue
            return content.strip()
        except Exception as e:
            logger.error(f"Error processing streaming response: {e}")
            raise RuntimeError(f"Error processing streaming response: {e}")

def create_agent(tools: List[Tool]) -> Any:
    """
    Initialize the LangChain agent with the provided tools.

    Args:
        tools (List[Tool]): List of LangChain Tool objects.

    Returns:
        Any: Initialized agent.
    """
    llm = LMStudioLLM()

    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=False,
        handle_parsing_errors=True,
    )
    return agent