import json
import logging
import re
import os
from groq import Groq
from tools import get_current_datetime, search_web, ReminderManager

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

class TelegramJadeAgent:
    def __init__(self, reminder_manager=None, config_path="telegram_jade/config.json"):
        # Simple config loading
        self.config = {}
        if os.path.exists(config_path):
            with open(config_path) as f:
                self.config = json.load(f)
        
        self.model = self.config.get("groq_model", "llama3-70b-8192")
        self.client = Groq(api_key=os.getenv("GROQ_API_KEY"))
        
        self.reminder_manager = reminder_manager if reminder_manager else ReminderManager()
        
        # We keep history per chat_id in memory for now (could be persisted later)
        self.histories = {}

    def _get_system_prompt(self):
        current_time = get_current_datetime()
        return f"""You are Jade, a 24h personal assistant on Telegram.
Current Time: {current_time}

Your goal is to be helpful, manage reminders, and search the web when needed.
You have a personality: friendly, efficient, and slightly witty.

AVAILABLE TOOLS:
You can call tools by responding ONLY with a JSON block in this format:
{{"tool": "tool_name", "args": {{"arg1": "value1"}}}}

Tools:
1. get_current_datetime() -> Returns current date/time.
2. search_web(query: str) -> Search internet for info.
3. schedule_reminder(time_str: str, message: str) -> Schedule a reminder. 
   - time_str MUST be in "YYYY-MM-DD HH:MM:SS" format. 
   - If the user gives a relative time (e.g., "in 10 mins"), CALCULATE the absolute time based on Current Time.
4. list_reminders() -> List pending reminders for the user.
5. delete_reminder(reminder_id: int) -> Delete a reminder.

RULES:
- Always check the Current Time before scheduling.
- If the user asks for a reminder, you MUST calculate the exact "YYYY-MM-DD HH:MM:SS" and use `schedule_reminder`.
- If you need to answer a question about recent events, use `search_web`.
- If you are just chatting, reply with plain text.
"""

    def _get_history(self, chat_id):
        if chat_id not in self.histories:
            self.histories[chat_id] = [{"role": "system", "content": self._get_system_prompt()}]
        else:
            # Update system prompt with fresh time
            self.histories[chat_id][0]["content"] = self._get_system_prompt()
        return self.histories[chat_id]

    def _process_tool_call(self, response_text):
        try:
            match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if not match:
                return None
            data = json.loads(match.group(0))
            if "tool" in data and "args" in data:
                return data
        except:
            pass
        return None

    def _run_tool(self, tool_data, chat_id):
        name = tool_data["tool"]
        args = tool_data["args"]
        
        logger.info(f"Executing tool {name} for chat {chat_id}")

        if name == "get_current_datetime":
            return get_current_datetime()
        elif name == "search_web":
            return search_web(args.get("query"))
        elif name == "schedule_reminder":
            # We need to hook this into the actual scheduler in main.py.
            # Ideally, the Agent should return this intent to the Main loop, 
            # or the ReminderManager handles the DB and Main loop watches the DB.
            # For simplicity: Agent updates DB via ReminderManager. 
            # Main loop (Scheduler) should refresh or listen to changes.
            # BUT: `schedule_reminder` here just updates the JSON. 
            # The Scheduler in `main.py` needs to be aware of new jobs.
            # We will return a special signal or just update the DB and assume Main reloads or we return a callback result.
            return self.reminder_manager.add_reminder(chat_id, args.get("time_str"), args.get("message"))
        elif name == "list_reminders":
            return self.reminder_manager.list_reminders(chat_id)
        elif name == "delete_reminder":
            return self.reminder_manager.delete_reminder(chat_id, int(args.get("reminder_id")))
        else:
            return f"Unknown tool: {name}"

    def chat(self, chat_id, user_input):
        history = self._get_history(chat_id)
        history.append({"role": "user", "content": user_input})
        
        # Simple ReAct loop
        for _ in range(5):
            try:
                completion = self.client.chat.completions.create(
                    messages=history,
                    model=self.model,
                    temperature=0.5
                )
                response = completion.choices[0].message.content
            except Exception as e:
                return f"Error calling Groq: {e}"

            tool_data = self._process_tool_call(response)
            
            if tool_data:
                history.append({"role": "assistant", "content": response})
                
                tool_result = self._run_tool(tool_data, chat_id)
                
                history.append({"role": "system", "content": f"TOOL_RESULT: {tool_result}"})
                
                # If it was a reminder scheduling, we might want to inform the caller (main.py) 
                # to update the scheduler.
                # For now, we just continue the conversation loop.
            else:
                history.append({"role": "assistant", "content": response})
                
                # Limit history size
                if len(history) > 20:
                    history = [history[0]] + history[-19:]
                self.histories[chat_id] = history
                
                return response
                
        return "I'm getting confused. Let's stop here."

    def generate_reminder_message(self, chat_id, reminder_message):
        """
        Called when a scheduled reminder triggers. 
        Generates a friendly notification message.
        """
        prompt = [
            {"role": "system", "content": "You are Jade. It is time to remind the user of something."},
            {"role": "user", "content": f"The reminder is: '{reminder_message}'. Write a friendly message to send to the user now."}
        ]
        try:
            completion = self.client.chat.completions.create(
                messages=prompt,
                model=self.model,
                temperature=0.7
            )
            return completion.choices[0].message.content
        except Exception as e:
            return f"Reminder: {reminder_message}"