main.py

import os
import baml_client
from baml_client.types import ProjectInput, Task, Schedule
from baml_client.async_client import b
from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Optional
import asyncio
from dotenv import load_dotenv


# Import the save function from our db module
from db import save_plan, get_db_connection
# Add these imports to the top of main.py
from datetime import datetime, timedelta
import calendar_integration

# Load environment variables (optional, if you have a .env file)
load_dotenv()

# --- Langgraph State Definition ---
class AppState(TypedDict):
    project_input: ProjectInput
    broken_down_tasks: Optional[List[Task]]
    schedule: Optional[Schedule]
    error: Optional[str]


# Add this function to main.py
async def get_calendar_constraints(days_ahead=7):
    """Get constraints from Google Calendar for the next week."""
    try:
        # Calculate date range (next 7 days by default)
        start_date = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
        end_date = start_date + timedelta(days=days_ahead)
        
        # Get events from Google Calendar
        events = calendar_integration.get_existing_events(start_date, end_date)
        
        # Format events as constraints
        cal_constraints = calendar_integration.format_events_for_constraints(events)
        
        return cal_constraints
    except Exception as e:
        print(f"Error getting calendar constraints: {e}")
        return []
    
# --- Langgraph Nodes ---
async def breakdown_node(state: AppState):
    print("--- Running Task Breakdown Node ---")
    project_input = state["project_input"]
    try:
        response = await b.BreakdownProject(input=project_input)
        print(
            "Breakdown successful:",
            [t.description for t in response.tasks]
            if response.tasks
            else "No tasks returned",
        )
        return {"broken_down_tasks": response.tasks}
    except Exception as e:
        print(f"Error during task breakdown: {e}")
        return {"error": f"BAML BreakdownProject failed: {e}"}


async def schedule_node(state: AppState):
    print("--- Running Scheduling Node ---")
    tasks = state.get("broken_down_tasks")
    constraints = state["project_input"].constraints
    if not tasks:
        print("No tasks found to schedule.")
        return {"error": "No tasks available from breakdown step."}

    try:
        schedule_response = await b.GenerateSchedule(
            tasks=tasks, constraints=constraints
        )
        print("Scheduling successful.")
        # Print scheduled tasks for clarity
        if schedule_response.scheduled_tasks:
            print("Scheduled Tasks:")
            for st in schedule_response.scheduled_tasks:
                print(f" - {st.task.description} ({st.start_time} to {st.end_time})")
        if schedule_response.unscheduled_tasks:
            print("Unscheduled Tasks:")
            for ut in schedule_response.unscheduled_tasks:
                print(f" - {ut.description}")

        return {"schedule": schedule_response}
    except Exception as e:
        print(f"Error during scheduling: {e}")
        return {"error": f"BAML GenerateSchedule failed: {e}"}


# --- Langgraph Graph Definition ---
print("Compiling Langgraph workflow...")
workflow = StateGraph(AppState)

workflow.add_node("breakdown", breakdown_node)
workflow.add_node("generate_schedule", schedule_node)

workflow.set_entry_point("breakdown")
workflow.add_edge("breakdown", "generate_schedule")
workflow.add_edge("generate_schedule", END)

app = workflow.compile()
print("Langgraph workflow compiled.")

async def run_agent(project_input: ProjectInput) -> AppState:
    """Runs the Langgraph agent with the given project input and saves the result."""
    # Try to get calendar constraints if credentials exist
    try:
        if calendar_integration.has_valid_credentials():
            calendar_constraints = await get_calendar_constraints()
            if calendar_constraints:
                # Add calendar constraints to the user's constraints
                project_input.constraints.extend(calendar_constraints)
                print(f"Added {len(calendar_constraints)} constraints from Google Calendar")
    except Exception as e:
        print(f"Error incorporating calendar constraints: {e}")
    
    initial_state = {"project_input": project_input}
    print(f"--- Running Agent for Goal: {project_input.goal} ---")
    final_state = await app.ainvoke(initial_state)
    
    # Save the result to the database if we have a schedule
    if final_state.get("schedule") and not final_state.get("error"):
        try:
            save_plan(project_input, final_state["schedule"])
            print("Plan saved to database.")
        except Exception as e:
            print(f"Error saving plan to database: {e}")
    
    return final_state

# --- Main Execution (for standalone testing) ---
async def main():
    # Example Input
    example_project_input = ProjectInput(
        goal="Write a blog post about AI productivity tools",
        existing_tasks=["Research existing tools", "Finalize outline"],
        constraints=[
            "Working hours: Mon-Fri 9am-5pm PST",
            "Lunch break: 12pm-1pm PST",
            "Deadline: End of day Friday",
            "Focus block: Tue 10am-12pm for writing",
        ],
    )

    print("--- Starting Standalone Test Run --- replicating second brain AI --- ")
    final_state = await run_agent(example_project_input)
    print("--- Standalone Test Run Finished ---")

    if final_state.get("error"):
        print("\n--- Errors ---")
        print(final_state["error"])
    elif final_state.get("schedule"):
        print(
            "\n--- Final Schedule (Standalone Test) --- replicating second brain AI --- "
        )
        # Schedule is already printed within the node, just confirm completion
        print("Schedule generated successfully.")
        # Optionally print the raw schedule object:
        # from pprint import pprint
        # pprint(final_state['schedule'].model_dump())
    else:
        print(
            "\n--- Workflow completed but no schedule generated or error reported. --- replicating second brain AI ---"
        )
        print(final_state)


if __name__ == "__main__":
    # Set BAML log level if needed (e.g., 'DEBUG', 'INFO', 'WARNING')
    # os.environ["BAML_LOG_LEVEL"] = "INFO"

    # Initialize DB connection explicitly when running standalone test
    # to ensure schema is created before running agent
    print("Initializing DB for standalone run...")
    get_db_connection()
    print("DB initialized.")

    asyncio.run(main())