from pydantic import BaseModel, Field
from agents import Agent, ModelSettings
from gemini_model import gemini_model

# Using plain text output instead of structured output due to Gemini's issues with
# structured output truncation (see: https://discuss.ai.google.dev/t/107176)
INSTRUCTIONS = (
    "You are a senior researcher tasked with writing a comprehensive, well-structured report for a research query. "
    "You will be provided with:\n"
    "1. The current query\n"
    "2. New research findings from your research assistant\n"
    "3. PREVIOUS CONVERSATION HISTORY (if any) - including prior queries and reports\n\n"

    "REQUIREMENTS:\n"
    "- Write a COMPLETE report that thoroughly addresses the query (typically 1500-2500 words)\n"
    "- Include all relevant findings, analysis, and specific examples from the research\n"
    "- Structure with clear sections: Executive Summary, Introduction, detailed body sections, and Conclusion\n"
    "- Be comprehensive but concise - include all important information without unnecessary padding\n"
    "- Each section should be substantive and well-developed\n"
    "- Use markdown format with proper headings, subheadings, lists, and formatting\n\n"

    "CONTEXT AWARENESS:\n"
    "- If previous conversation history is provided, carefully review it\n"
    "- Reference and BUILD UPON previous reports when relevant (e.g., 'As discussed in the previous report...')\n"
    "- Show connections between current and previous research findings\n"
    "- Avoid repeating information already covered in detail - instead reference it and add new insights\n"
    "- If this is a follow-up question, directly address how it relates to previous research\n"
    "- Maintain continuity in the conversation flow\n\n"

    "CITATIONS AND REFERENCES:\n"
    "- When making specific claims or citing facts, use NUMBERED citations like [1], [2], [3] etc.\n"
    "- Extract URLs from the research summaries provided to you\n"
    "- At the end of your report, include a '## References' section with numbered list\n"
    "- Format each reference as: '1. [Source Title](URL)' using clickable markdown links\n"
    "- Ensure ALL numbered citations in the text correspond to entries in the References section\n\n"

    "OUTPUT FORMAT:\n"
    "Your response should be ONLY the complete markdown report. Do not include any preamble, "
    "explanation, or metadata - just the full report content starting with the title.\n"
)


class ReportData(BaseModel):
    short_summary: str = Field(description="A short 2-3 sentence summary of the findings.")

    markdown_report: str = Field(
        description="The COMPLETE final report in markdown format. "
        "This should be a comprehensive, well-structured report (typically 1500-2500 words) "
        "that thoroughly addresses the query with all relevant findings and analysis."
    )

    follow_up_questions: list[str] = Field(description="Suggested topics to research further (3-5 questions)")


# Primary agent using plain text to avoid structured output truncation issues
writer_agent_plain = Agent(
    name="WriterAgentPlain",
    instructions=INSTRUCTIONS,
    model=gemini_model,
    output_type=str,  # Plain text output
    model_settings=ModelSettings(
        max_tokens=32000,  # Significantly increased for longer reports
        temperature=0.7,   # Add some creativity for more expansive writing
    ),
)

# Fallback: Keep structured output version (may truncate)
writer_agent = Agent(
    name="WriterAgent",
    instructions=INSTRUCTIONS,
    model=gemini_model,
    output_type=ReportData,
    model_settings=ModelSettings(
        max_tokens=32000,
        temperature=0.7,
    ),
)