import os
import asyncio
import argparse
from pathlib import Path

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# -------------------- 1. LLM and Embeddings --------------------
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,
)

embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
)

# -------------------- 2. Chroma DB --------------------
CHROMA_PATH = Path("./chroma_faq")
CHROMA_COLLECTION = "course_faq"

vector_store = Chroma(
    collection_name=CHROMA_COLLECTION,
    embedding_function=embeddings,
    persist_directory=str(CHROMA_PATH),
)

# Persist changes
vector_store.persist()

# -------------------- 3. Tools --------------------
@tool
def search_knowledge(query: str) -> str:
    """Search the knowledge base for relevant information."""
    docs = vector_store.similarity_search(query, k=3)
    return "\n".join(d.page_content for d in docs) if docs else "No results found in course materials."

@tool
def fetch_course_meta(query: str) -> str:
    """Fetch course metadata (schedule, syllabus, etc.) from a static JSON file."""
    meta_path = Path("meta.json")
    if not meta_path.exists():
        return "Metadata file not found."
    import json
    data = json.loads(meta_path.read_text())
    query_lower = query.lower()
    if "schedule" in query_lower:
        return f"Course schedule: {data.get('schedule', 'Not available')}"
    if "syllabus" in query_lower:
        return f"Syllabus URL: {data.get('syllabus_url', 'Not available')}"
    if "instructor" in query_lower:
        return f"Instructor: {data.get('instructor', 'Not available')}"
    # Default: return all
    return json.dumps(data, indent=2)

# -------------------- 4. Backend --------------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# -------------------- 5. Agent --------------------
SYSTEM_PROMPT = (
    "You are a helpful FAQ bot for the course.\n"
    "Use the 'search_knowledge' tool to answer questions about course materials.\n"
    "Use the 'fetch_course_meta' tool for questions about schedule, syllabus, instructor, etc.\n"
    "Do not call both tools unnecessarily.\n"
    "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate which tool provided the information."
)

agent = create_deep_agent(
    model=llm,
    tools=[search_knowledge, fetch_course_meta],
    backend=backend,
    system_prompt=SYSTEM_PROMPT,
)

# -------------------- 6. Data Loader --------------------

def load_faq_to_chroma():
    data_dir = Path("data")
    if not data_dir.exists():
        data_dir.mkdir()
        # Create sample markdown files if missing
        (data_dir / "file1.md").write_text(
            "## Course Overview\nThis course covers advanced topics in AI. Topics include machine learning, deep learning, and natural language processing."
        )
        (data_dir / "file2.md").write_text(
            "## FAQ\nQ: What is the schedule?\nA: Sessions are held on Mondays and Wednesdays.\nQ: Where can I find the syllabus?\nA: Syllabus is available on the course website."
        )
    docs = []
    for md_file in data_dir.glob("*.md"):
        text = md_file.read_text()
        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
    vector_store.add_documents(docs)
    vector_store.persist()
    print(f"Loaded {len(docs)} documents into Chroma collection '{CHROMA_COLLECTION}'.")

# -------------------- 7. CLI --------------------
PRESET_QUESTIONS = {
    "1": "What topics are covered in the course?",
    "2": "Where can I find the syllabus?",
    "3": "What is the course schedule?",
}

async def run_question(question: str, thread_id: str = "session-1"):
    result = await agent.ainvoke(
        {"messages": ["HumanMessage(content=\"{}\")".format(question)]},
        {"configurable": {"thread_id": thread_id}},
    )
    # The agent returns a dict with 'messages'; extract last message
    content = result["messages"][-1].content
    print(content)

async def main():
    # Ensure data is loaded
    load_faq_to_chroma()

    parser = argparse.ArgumentParser(description="FAQ Bot CLI")
    parser.add_argument("--preset", choices=["1", "2", "3"], help="Run a preset question")
    args = parser.parse_args()

    if args.preset:
        question = PRESET_QUESTIONS[args.preset]
        print(f"Preset question {args.preset}: {question}")
        await run_question(question)
    else:
        print("Enter your question (type 'exit' to quit):")
        while True:
            q = input("> ")
            if q.lower() in {"exit", "quit"}:
                break
            if q.strip():
                await run_question(q)

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