import os
import json
import logging
from datetime import datetime, timedelta
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.schema import SystemMessage, HumanMessage

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

class Agent:
    def __init__(self, role: str, goal: str, backstory: str, personality: str = "", llm=None) -> None:
        """
        Initialize an Agent with role, goal, backstory, personality, and assigned LLM.
        """
        self.role = role
        self.goal = goal
        self.backstory = backstory
        self.personality = personality
        self.tools = []  # Initialize with empty list for future tool integrations
        self.llm = llm

class Task:
    def __init__(self, description: str, agent: Agent, expected_output: str, context=None) -> None:
        """
        Initialize a Task with its description, the responsible agent, expected output, and optional context.
        """
        self.description = description
        self.agent = agent
        self.expected_output = expected_output
        self.context = context or []

google_api_key = os.getenv("GEMINI_API_KEY")  # 실제 Google API 키 사용
if not google_api_key:
    logging.error("GEMINI_API_KEY is not set in the environment variables.")
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=google_api_key)

# -------------------------------------------------------------------------------
# Define Academic Research Agents
# -------------------------------------------------------------------------------
literature_research_agent = Agent(
    role="Literature Research Agent",
    goal="Research and provide a comprehensive review of existing literature on the research topic.",
    backstory="An experienced academic researcher specialized in literature reviews and meta-analyses.",
    personality="Analytical, thorough, and detail-oriented.",
    llm=llm,
)

outline_agent = Agent(
    role="Outline Agent",
    goal="Generate a structured and detailed outline for a research paper based on the research topic and literature.",
    backstory="A methodical academic planner who organizes research findings into coherent paper structures.",
    personality="Organized, systematic, and insightful.",
    llm=llm,
)

draft_writing_agent = Agent(
    role="Draft Writing Agent",
    goal="Compose a first draft of the research paper based on the literature review and outline.",
    backstory="A skilled academic writer capable of synthesizing research findings into well-structured drafts.",
    personality="Articulate, precise, and scholarly.",
    llm=llm,
)

citation_agent = Agent(
    role="Citation Agent",
    goal="Generate a list of relevant citations and references in the required format for the research paper.",
    backstory="A detail-oriented bibliographic expert with extensive knowledge of citation standards.",
    personality="Meticulous, accurate, and research-savvy.",
    llm=llm,
)

editing_agent = Agent(
    role="Editing Agent",
    goal="Revise and polish the draft for clarity, coherence, and academic tone.",
    backstory="An expert editor skilled in improving academic manuscripts and ensuring high-quality presentation.",
    personality="Critical, precise, and supportive.",
    llm=llm,
)

chatbot_agent = Agent(
    role="Chatbot Agent",
    goal="Engage in interactive conversation to answer queries related to the academic research process.",
    backstory="A conversational AI assistant with extensive knowledge in academia and research methodologies.",
    personality="Helpful, conversational, and knowledgeable.",
    llm=llm,
)

# -------------------------------------------------------------------------------
# Define Tasks for Academic Research and Writing
# -------------------------------------------------------------------------------
literature_research_task = Task(
    description="""Research academic literature on {topic} considering the keywords {keywords}.
    
Please provide:
- A summary of the current state of research,
- Key trends and gaps in the literature,
- Notable studies and their findings,
- Relevant theoretical frameworks and methodologies.
Format the response with bullet points and concise summaries.""",
    agent=literature_research_agent,
    expected_output="""A comprehensive literature review summary covering:
1. Summary of current research trends
2. Identification of gaps and controversies
3. Key studies with brief descriptions
4. Theoretical frameworks and methodologies used"""
)

outline_task = Task(
    description="""Based on the research topic {topic} and literature review findings, generate a detailed outline for a research paper.
    
Include sections such as:
- Abstract
- Introduction (including research questions and objectives)
- Literature Review
- Methodology
- Results/Findings
- Discussion
- Conclusion
- References
Format the outline in a structured manner with bullet points and subheadings.""",
    agent=outline_agent,
    expected_output="A structured outline for a research paper including all major sections and key points to cover in each section."
)

draft_writing_task = Task(
    description="""Using the research topic {topic}, the literature review, and the generated outline, compose a first draft of the research paper.
    
The draft should include:
- A coherent narrative flow,
- Detailed sections as per the outline,
- Integration of key findings from the literature review.
Ensure the tone is academic and the content is well-organized.""",
    agent=draft_writing_agent,
    expected_output="A complete first draft of the research paper covering all sections with sufficient academic detail."
)

citation_task = Task(
    description="""Based on the literature review for {topic}, generate a list of key references and citations in APA format.
    
Include:
- Author names, publication year, title, and source,
- At least 10 key references relevant to the research topic.
Format the output as a numbered list of citations.""",
    agent=citation_agent,
    expected_output="A list of 10+ relevant citations in APA format."
)

editing_task = Task(
    description="""Review and edit the draft for clarity, coherence, and academic tone.
    
Focus on:
- Improving sentence structure,
- Ensuring logical flow between sections,
- Correcting grammar and stylistic issues,
- Enhancing academic tone.
Provide the polished version of the paper.""",
    agent=editing_agent,
    expected_output="A refined and polished version of the research paper draft."
)

chatbot_task = Task(
    description="Provide a conversational and detailed response to academic research-related queries.",
    agent=chatbot_agent,
    expected_output="A friendly, informative response addressing the query."
)

def run_task(task: Task, input_text: str) -> str:
    """
    Executes the given task using the associated agent's LLM and returns the response content.
    """
    try:
        if not isinstance(task, Task):
            raise ValueError(f"Expected 'task' to be an instance of Task, got {type(task)}")
        if not hasattr(task, 'agent') or not isinstance(task.agent, Agent):
            raise ValueError("Task must have a valid 'agent' attribute of type Agent.")
        system_input = (
            f"Agent Details:\n"
            f"Role: {task.agent.role}\n"
            f"Goal: {task.agent.goal}\n"
            f"Backstory: {task.agent.backstory}\n"
            f"Personality: {task.agent.personality}\n"
        )
        task_input = (
            f"Task Details:\n"
            f"Task Description: {task.description}\n"
            f"Expected Output: {task.expected_output}\n"
            f"Input for Task:\n{input_text}\n"
        )
        messages = [
            SystemMessage(content=system_input),
            HumanMessage(content=task_input)
        ]
        response = task.agent.llm.invoke(messages)
        if not response or not response.content:
            raise ValueError("Empty response from LLM.")
        return response.content
    except Exception as e:
        logging.error(f"Error in task '{task.agent.role}': {e}")
        return f"Error in {task.agent.role}: {e}"