import os
import json
import logging
from datetime import datetime, timedelta
from langchain_groq import ChatGroq
from langchain.schema import SystemMessage, HumanMessage
from dotenv import load_dotenv
load_dotenv()

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:
        self.role = role
        self.goal = goal
        self.backstory = backstory
        self.personality = personality
        self.tools = []
        self.llm = llm

class Task:
    def __init__(self, description: str, agent: Agent, expected_output: str, context=None) -> None:
        self.description = description
        self.agent = agent
        self.expected_output = expected_output
        self.context = context or []

groq_api_key = os.getenv("GROQ_API_KEY")
os.environ['GROQ_API_KEY'] = groq_api_key
if not groq_api_key:
    logging.error("GROQ_API_KEY is not set in the environment variables.")
llm = ChatGroq(model="llama-3.3-70b-versatile", max_tokens=4000)

literature_research_agent = Agent(
    role="Literature Research Agent",
    goal="Conduct an in-depth review of existing scholarly work relevant to a given research theme.",
    backstory="A seasoned academic investigator skilled at uncovering trends and gaps in published research.",
    personality="Analytical, detail-focused, methodical.",
    llm=llm,
)

outline_agent = Agent(
    role="Outline Agent",
    goal="Design a logical, in-depth framework for organizing academic content effectively.",
    backstory="An academic strategist who formulates detailed paper structures aligned with scholarly standards.",
    personality="Structured, organized, and perceptive.",
    llm=llm,
)

draft_writing_agent = Agent(
    role="Draft Writing Agent",
    goal="Generate a refined academic draft integrating conceptual and empirical elements seamlessly.",
    backstory="An articulate academic writer who constructs comprehensive and original content based on guidance.",
    personality="Clear, scholarly, and refined.",
    llm=llm,
)

citation_agent = Agent(
    role="Citation Agent",
    goal="Curate precise bibliographic references aligned with formatting norms.",
    backstory="A citation expert well-versed in academic documentation and referencing styles.",
    personality="Exacting, thorough, and standards-oriented.",
    llm=llm,
)

editing_agent = Agent(
    role="Editing Agent",
    goal="Enhance clarity, style, and scholarly tone while preserving originality and academic rigor.",
    backstory="A language and style expert who ensures manuscripts are professional and naturally flowing.",
    personality="Discerning, articulate, and attentive.",
    llm=llm,
)

chatbot_agent = Agent(
    role="Chatbot Agent",
    goal="Answer academic inquiries and assist with research-related tasks in a conversational manner.",
    backstory="A helpful academic consultant embedded in an AI system to support researchers.",
    personality="Conversational, knowledgeable, and responsive.",
    llm=llm,
)

literature_research_task = Task(
    description="""Conduct a contextual review of academic studies pertaining to {topic}, integrating relevant concepts, trends, and methodological approaches using the keywords {keywords}.

Avoid generic phrasing and generate a unique articulation by:
- Summarizing key themes across disciplines,
- Highlighting research voids and disputed findings,
- Referring to pivotal studies with paraphrased insights,
- Discussing theoretical models and research paradigms involved.
Use informative bullet points and neutral scholarly tone.
    """,
    agent=literature_research_agent,
    expected_output="""A well-contextualized literature summary that includes:
1. Evolving thematic patterns
2. Unresolved issues and gaps
3. Specific empirical works (paraphrased)
4. Conceptual lenses and methodologies used
    """
)

outline_task = Task(
    description="""Based on the topic {topic} and derived insights from prior literature, outline a robust research paper structure that reflects academic conventions.

Organize into:
- Concise Abstract
- Framing Introduction (with aims/questions)
- Synthesized Literature Review
- Transparent Methodology
- Expected or actual Results
- Interpretative Discussion
- Summative Conclusion
- Citations placeholder
Use bullet points and avoid repeating standard textbook language.
    """,
    agent=outline_agent,
    expected_output="A coherent research paper outline formatted with headings, subheadings, and key elements in each part."
)

draft_writing_task = Task(
    description="""Utilizing the theme {topic}, literature review, and proposed structure, compose an initial draft of an academic paper.

Ensure:
- Flow of argumentation is natural,
- Terminology varies and mimics human phrasing,
- Integration of insights without obvious reuse of source phrasing,
- Academic tone that appears manually crafted.

Keep the prose clear but layered, avoiding repetition or formulaic transitions.
    """,
    agent=draft_writing_agent,
    expected_output="An academically styled, logically progressive draft that reads as if independently composed."
)

citation_task = Task(
    description="""Produce APA-style references for scholarly works relating to {topic}, drawing from domain-appropriate journals and books.

Use human-like variability in citation content and structure:
- List 10 or more significant sources,
- Vary author name formats slightly,
- Paraphrase titles subtly,
- Represent realistic and domain-appropriate publishers or journals.
Output a numbered list with natural inconsistencies typical of human formatting.
    """,
    agent=citation_agent,
    expected_output="A diverse and natural-feeling list of scholarly APA references."
)

editing_task = Task(
    description="""Perform a refined editorial pass over the draft, rewriting where necessary to make it appear distinct from LLM outputs.

Steps:
- Vary syntax and phrasing complexity,
- Insert nuanced transitions,
- Restructure paragraphs subtly,
- Remove indicators of AI authorship.
Output the fully edited and naturally flowing version.
    """,
    agent=editing_agent,
    expected_output="A polished, stylistically diverse version that avoids detectable patterns of machine-generated content."
)

chatbot_task = Task(
    description="""Engage in research-related academic discussions or answer questions while maintaining a tone that is informal but intellectually competent.

Ensure responses sound helpful, nuanced, and naturally composed.
    """,
    agent=chatbot_agent,
    expected_output="Conversational yet knowledgeable academic assistance."
)

def run_task(task: Task, input_text: str) -> str:
    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}"