import os
import asyncio
from typing import Union

from dotenv import load_dotenv
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
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
from deepagents.tools import tool

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 in years, optional if not mentioned"
    )
    profession: str = Field(description="Professional title or occupation")
    skills: list[str] = Field(description="List of key skills or technologies")


class MeetingNotes(BaseModel):
    date: str = Field(description="Date of the meeting in ISO format (YYYY-MM-DD)")
    participants: list[str] = Field(description="List of participant names")
    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 or next steps")


# ----------------------------------------------------------------------
# LLM configuration (OpenRouter via langchain-openai)
# ----------------------------------------------------------------------
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,
)

# ----------------------------------------------------------------------
# Prompt templates with format instructions
# ----------------------------------------------------------------------
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)

person_prompt = PromptTemplate.from_template(
    """Extract the following information about a person and output it as JSON that matches the given schema.

{format_instructions}

Text:
{input_text}
"""
)

meeting_prompt = PromptTemplate.from_template(
    """Extract structured meeting notes from the given text and output them as JSON that matches the given schema.

{format_instructions}

Text:
{input_text}
"""
)

# ----------------------------------------------------------------------
# Simple routing based on keyword heuristics
# ----------------------------------------------------------------------
def select_schema(text: str) -> str:
    """Return 'person' or 'meeting' depending on the content."""
    lowered = text.lower()
    meeting_keywords = ["встреча", "meeting", "участники", "participants", "agenda", "решения"]
    if any(word in lowered for word in meeting_keywords):
        return "meeting"
    return "person"


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

@tool
def route_and_process(text: str) -> str:
    """Detect the type of the input text, run the appropriate extraction chain,
    and return the JSON representation of the validated Pydantic model."""
    schema_type = select_schema(text)

    if schema_type == "person":
        parser = person_parser
        prompt = person_prompt
    else:
        parser = meeting_parser
        prompt = meeting_prompt

    chain = (
        prompt.partial(format_instructions=parser.get_format_instructions())
        | llm
        | parser
    )
    result = chain.invoke({"input_text": text})
    # Return pretty JSON for CLI display
    return result.model_dump_json(indent=2)


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


# ----------------------------------------------------------------------
# CLI interface
# ----------------------------------------------------------------------
EXAMPLES = {
    "person": "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.",
    "meeting": """Дата: 2024-09-15
Участники: Иван, Мария, Алексей
Темы: План проекта, бюджет, сроки
Решения: Утвердить бюджет в 500k, начать разработку 1 октября
Следующие шаги: Иван подготовит ТЗ, Мария соберёт требования, Алексей настроит окружение."""
}


async def run_example(example_key: str):
    text = EXAMPLES[example_key]
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=text)]},
        {"configurable": {"thread_id": f"example-{example_key}"}},
    )
    print("Input text:")
    print(text)
    print("\nExtracted JSON:")
    print(result["messages"][-1].content)


def interactive_mode():
    print("Enter text (empty line to finish):")
    lines = []
    while True:
        line = input()
        if line == "":
            break
        lines.append(line)
    user_text = "\n".join(lines)
    if not user_text.strip():
        print("No input provided.")
        return

    result = asyncio.run(
        agent.ainvoke(
            {"messages": [HumanMessage(content=user_text)]},
            {"configurable": {"thread_id": "interactive-session"}},
        )
    )
    print("\nExtracted JSON:")
    print(result["messages"][-1].content)


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Structured extraction demo")
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--example",
        choices=["person", "meeting"],
        help="Run a built-in example",
    )
    group.add_argument(
        "--interactive",
        action="store_true",
        help="Enter interactive mode",
    )
    args = parser.parse_args()

    if args.example:
        asyncio.run(run_example(args.example))
    elif args.interactive:
        interactive_mode()
    else:
        parser.print_help()


if __name__ == "__main__":
    main()