import os
import asyncio
import argparse
from dotenv import load_dotenv

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate

from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# Load API key from .env
load_dotenv()

# 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,
)

# Pydantic models
class PersonInfo(BaseModel):
    name: str = Field(description="Full name of the person")
    age: int | None = Field(description="Age of the person, optional", default=None)
    profession: str = Field(description="Current profession or role")
    skills: list[str] = Field(description="List of skills or technologies")

class MeetingNotes(BaseModel):
    date: str = Field(description="Date of the meeting (YYYY-MM-DD)")
    participants: list[str] = Field(description="Names of participants")
    topics: list[str] = Field(description="Main discussion topics")
    decisions: list[str] = Field(description="Decisions made during the meeting")
    next_steps: list[str] = Field(description="Action items for next steps")

# Prompt templates
person_prompt = PromptTemplate.from_template(
    """
Extract the following information from the text: name, age (optional), profession, skills.
The text: {input_text}
Output must be in JSON format according to the following schema:
{format_instructions}
"""
)

meeting_prompt = PromptTemplate.from_template(
    """
Extract the following information from the text: date, participants, topics, decisions, next_steps.
The text: {input_text}
Output must be in JSON format according to the following schema:
{format_instructions}
"""
)

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

# Simple heuristic to decide which schema to use

def determine_schema(text: str) -> str:
    lower = text.lower()
    meeting_keywords = ["meeting", "participants", "date", "topics", "decisions", "next steps", "next_step", "next_step"]
    if any(k in lower for k in meeting_keywords):
        return "meeting"
    return "person"

# Summary generation

def create_summary(parsed: BaseModel, schema: str) -> str:
    if schema == "meeting":
        return (
            f"Meeting on {parsed.date} with participants {', '.join(parsed.participants)}. "
            f"Topics: {', '.join(parsed.topics)}. Decisions: {', '.join(parsed.decisions)}. "
            f"Next steps: {', '.join(parsed.next_steps)}."
        )
    else:
        age = parsed.age if parsed.age is not None else "N/A"
        return (
            f"{parsed.name}, age {age}, profession {parsed.profession}. "
            f"Skills: {', '.join(parsed.skills)}."
        )

# Tool that performs extraction
@tool
def process_text(input_text: str) -> str:
    schema = determine_schema(input_text)
    if schema == "meeting":
        prompt = meeting_prompt
        parser = meeting_parser
    else:
        prompt = person_prompt
        parser = person_parser
    chain = prompt | llm | parser
    parsed = chain.invoke({"input_text": input_text})
    summary = create_summary(parsed, schema)
    return parsed.model_dump_json(indent=2) + "\n\nSummary: " + summary

# Backend for the agent
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# Create the deep agent
agent = create_deep_agent(
    model=llm,
    tools=[process_text],
    backend=backend,
    system_prompt="You are a helpful agent that extracts structured data from text.",
)

# CLI entry point
async def main():
    parser = argparse.ArgumentParser(description="Extract structured data from text.")
    parser.add_argument("--example", choices=["person", "meeting"], help="Run example")
    parser.add_argument("--text", type=str, help="Input text")
    args = parser.parse_args()

    if args.example:
        if args.example == "person":
            input_text = "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker."
        else:
            input_text = (
                "Дата: 2023-05-01. Участники: Иван, Мария. Темы: проект X, бюджет. "
                "Решения: утвердить бюджет. Next steps: подготовить план."
            )
    elif args.text:
        input_text = args.text
    else:
        input_text = input("Enter text: ")

    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=input_text)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    output = result["messages"][-1].content
    print(output)

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