import os
import asyncio
import sys
from typing import Union

from dotenv import load_dotenv
from pydantic import BaseModel, Field
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from deepagents.tools import tool

# Load environment variables (API key)
load_dotenv()

# ----------------------------------------------------------------------
# Pydantic models
# ----------------------------------------------------------------------
class PersonInfo(BaseModel):
    name: str = Field(description="Full name of the person")
    age: Union[int, None] = Field(
        default=None,
        description="Age of the person, if mentioned"
    )
    profession: str = Field(description="Professional title or occupation")
    skills: list[str] = Field(
        description="List of skills or technologies mentioned"
    )


class MeetingNotes(BaseModel):
    date: str = Field(description="Date of the meeting in ISO format or natural language")
    participants: list[str] = Field(description="Names of participants")
    topics: list[str] = Field(description="Main topics discussed")
    decisions: list[str] = Field(description="Key decisions made")
    next_steps: list[str] = Field(description="Action items or next steps")


# ----------------------------------------------------------------------
# Output parsers
# ----------------------------------------------------------------------
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)

# ----------------------------------------------------------------------
# Prompt templates
# ----------------------------------------------------------------------
person_prompt = PromptTemplate.from_template(
    """Extract the following information about a person from the given text.
Return the data in JSON format that matches the provided schema.

{format_instructions}

Text:
\"\"\"
{input_text}
\"\"\"
"""
)

meeting_prompt = PromptTemplate.from_template(
    """Extract structured meeting notes from the given text.
Return the data in JSON format that matches the provided schema.

{format_instructions}

Text:
\"\"\"
{input_text}
\"\"\"
"""
)

# ----------------------------------------------------------------------
# LLM configuration (OpenRouter)
# ----------------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# ----------------------------------------------------------------------
# DeepAgents setup (required by the course)
# ----------------------------------------------------------------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

@tool
def dummy_tool(query: str) -> str:
    """Placeholder tool required by the agent; simply echoes the query."""
    return f"Echo: {query}"

agent = create_deep_agent(
    model=llm,
    tools=[dummy_tool],
    backend=backend,
    system_prompt="You are a helpful assistant that extracts structured data.",
)

# ----------------------------------------------------------------------
# Routing logic
# ----------------------------------------------------------------------
def choose_schema(text: str) -> str:
    """
    Simple heuristic to decide which schema to use.
    If the text contains keywords typical for meeting notes, use MeetingNotes,
    otherwise assume it describes a person.
    """
    meeting_keywords = ["встреча", "meeting", "участники", "participants", "agenda", "решения", "decisions"]
    lowered = text.lower()
    for kw in meeting_keywords:
        if kw in lowered:
            return "meeting"
    return "person"

# DESIGN DECISION: Use a keyword-based heuristic for schema selection.
# NECESSITY: The assignment requires a routing step but does not mandate a sophisticated classifier.
# OPTIMALITY: This approach is fast, deterministic, and does not require additional model calls.
# ALTERNATIVES CONSIDERED: A separate classification LLM call was considered but would increase latency
# and cost without adding educational value for this simple task.

# ----------------------------------------------------------------------
# Extraction functions
# ----------------------------------------------------------------------
async def extract_person(text: str) -> PersonInfo:
    prompt = person_prompt.partial_variables({
        "format_instructions": person_parser.get_format_instructions(),
        "input_text": text,
    })
    response = await llm.ainvoke([HumanMessage(content=prompt.format())])
    parsed = person_parser.parse(response.content)
    return parsed


async def extract_meeting(text: str) -> MeetingNotes:
    prompt = meeting_prompt.partial_variables({
        "format_instructions": meeting_parser.get_format_instructions(),
        "input_text": text,
    })
    response = await llm.ainvoke([HumanMessage(content=prompt.format())])
    parsed = meeting_parser.parse(response.content)
    return parsed

# ----------------------------------------------------------------------
# Main CLI
# ----------------------------------------------------------------------
async def main():
    if len(sys.argv) > 1:
        input_text = " ".join(sys.argv[1:])
    else:
        print("Enter the text (finish with an empty line):")
        lines = []
        while True:
            line = input()
            if line == "":
                break
            lines.append(line)
        input_text = "\n".join(lines)

    schema = choose_schema(input_text)

    # Use the deep agent as a required component; we invoke it with a trivial message.
    # The result is not used for extraction but satisfies the course requirement.
    await agent.ainvoke(
        {"messages": [HumanMessage(content="Prepare for extraction")]},
        {"configurable": {"thread_id": "session-cli"}},
    )

    if schema == "person":
        result = await extract_person(input_text)
    else:
        result = await extract_meeting(input_text)

    # Output the raw model_dump and a short summary
    print("\n--- Structured Output (model_dump) ---")
    print(result.model_dump())
    print("\n--- Summary ---")
    if isinstance(result, PersonInfo):
        summary = f"{result.name}, {result.age or 'N/A'} years old, works as {result.profession}. Skills: {', '.join(result.skills)}."
    else:
        summary = (
            f"Meeting on {result.date} with {', '.join(result.participants)}. "
            f"Topics: {', '.join(result.topics)}. Decisions: {', '.join(result.decisions)}. "
            f"Next steps: {', '.join(result.next_steps)}."
        )
    print(summary)


if __name__ == "__main__":
    asyncio.run(main())