import json
import re
from datetime import datetime
import pytz
from src.agents.base import BaseAgent
from src.core.agent_middleware import AgentMiddleware

class ManagerAgent(BaseAgent):
    def __init__(self, engine, memory, kb=None):
        super().__init__(engine, "manager")
        self.memory = memory
        self.middleware = AgentMiddleware()

    def analyze_task(self, task: str, history_str: str = ""):
        task_lower = task.lower().strip()
        if any(x in task_lower for x in ["viết bài", "quảng cáo", "content"]): return {"category": "MARKETING"}
        if any(x in task_lower for x in ["tạo", "build", "automation", "quy trình", "workflow"]): return {"category": "TECHNICAL"}
        if any(x in task_lower for x in ["doanh thu", "tồn kho", "bán được", "sales"]): return {"category": "DATA_INTERNAL"}
        return {"category": "GENERAL"}

    def plan(self, task: str, history_str: str = "", store_context: str = ""):
        schema = self.middleware.get_db_schema()
        prompt = f'''<|im_start|>system
You are the Architect. Design a logical flow for this automation request.
Available Data Schema:
{schema}
OUTPUT: A step-by-step logical plan.
<|im_end|>
<|im_start|>user
REQUEST: {task}
<|im_end|>
<|im_start|>assistant
'''
        return self.generate(prompt, max_new_tokens=512)

    def consult(self, task: str, context_data: str = "", history_str: str = "", store_context: str = ""):
        schema = self.middleware.get_db_schema()
        prompt = f'''<|im_start|>system
You are Project A, a Retail Assistant.
DATABASE SCHEMA:
{schema}
HISTORY:
{history_str}
DATA: {context_data}
INSTRUCTION: Answer helpfully in Vietnamese.
<|im_end|>
<|im_start|>user
{task}
<|im_end|>
<|im_start|>assistant
'''
        return self.generate(prompt, max_new_tokens=1024)
    
    def write_marketing(self, task: str):
        return self.generate(f"<|im_start|>system\nCopywriter.\n<|im_end|>\n<|im_start|>user\n{task}<|im_end|>\n<|im_start|>assistant\n", max_new_tokens=1024)