import uuid
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from langchain_core.messages import BaseMessage, HumanMessage, trim_messages
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from pydantic import BaseModel
from typing import Optional
import json
from sse_starlette.sse import EventSourceResponse
from datetime import datetime
from fastapi import APIRouter
from langchain_core.runnables import RunnableConfig
from langchain_core.prompts import ChatPromptTemplate
from typing import Any

router = APIRouter(
    prefix="/presentation",
    tags=["presentation"]
)

import json

@tool(response_format="content_and_artifact")
def plan(slides_json: str) -> tuple[str, dict]:
    """Create a presentation plan from a JSON string of slides (keys=slide numbers, values=content)."""
    try:
        slides = json.loads(slides_json)
        print(slides)
        return (
            f"Plan created with {len(slides)} slides: {', '.join(slides.keys())}.",
            {"slides_plan_json": slides_json}
        )
    except Exception as e:
        return (
            f"Invalid JSON format. Please provide a valid JSON string {str(e)[:100]}.",
            None
        )

@tool(response_format="content_and_artifact") 
def create_slide(slide_number: int, content: str, config: RunnableConfig) -> tuple[str, dict]:
    """Create a slide with the given number and content."""
    # Integration with slide creation API or template would go here
    slide = {
        "number": slide_number,
        "content": content,
        "created_at": datetime.now().isoformat()
    }
    return (
        f"Slide {slide_number} created",
        {"slide": slide}
    )

@tool(parse_docstring=True)
def execute_python(expression: str) -> str:
    """Execute a python mathematic expression. Returns the result of the expression or an error message if execution fails.

    Args:
        expression: The python expression to execute.
    """
    try:
        result = eval(expression)
        return f"The result of the expression is {result}"
    except Exception as e:
        return f"Error executing the expression: {str(e)}"

memory = MemorySaver()
model = ChatOpenAI(model="gpt-4o-mini", streaming=True)
prompt = ChatPromptTemplate.from_messages([
    ("system", """You are a Presentation Creation Assistant. Your task is to help users create effective presentations.
    Follow these steps:
    1. First use the plan tool to create an outline of the presentation
    2. Then use create_slide tool for each slide in sequence
    3. Guide the user through the presentation creation process
    
    Today's date is {{datetime.now().strftime('%Y-%m-%d')}}"""),
    ("placeholder", "{messages}"),
])

def state_modifier(state) -> list[BaseMessage]:
    try:
        formatted_prompt = prompt.invoke({
            "messages": state["messages"]
        })
        return trim_messages(
            formatted_prompt,
            token_counter=len,
            max_tokens=16000,
            strategy="last",
            start_on="human",
            include_system=True,
            allow_partial=False,
        )
    except Exception as e:
        print(f"Error in state modifier: {str(e)}")
        return state["messages"]

# Create the agent with presentation tools
agent = create_react_agent(
    model,
    tools=[plan, create_slide, execute_python],
    checkpointer=memory,
    state_modifier=state_modifier,
)

class ChatInput(BaseModel):
    message: str
    thread_id: Optional[str] = None

@router.post("/chat")
async def chat(input_data: ChatInput):
    thread_id = input_data.thread_id or str(uuid.uuid4())
    
    config = {
        "configurable": {
            "thread_id": thread_id
        }
    }
    
    input_message = HumanMessage(content=input_data.message)
    
    async def generate():
        async for event in agent.astream_events(
            {"messages": [input_message]}, 
            config,
            version="v2"
        ):
            kind = event["event"]
            
            if kind == "on_chat_model_stream":
                content = event["data"]["chunk"].content
                if content:
                    yield f"{json.dumps({'type': 'token', 'content': content})}\n"

            elif kind == "on_tool_start":
                tool_input = str(event['data'].get('input', ''))
                yield f"{json.dumps({'type': 'tool_start', 'tool': event['name'], 'input': tool_input})}\n"
            
            elif kind == "on_tool_end":
                print(event['data'])
                tool_output = event['data'].get('output', '')
                artifact_output = tool_output.artifact if tool_output.artifact else None
                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output.pretty_repr(), 'artifacts_data': artifact_output})}\n"
                print(tool_output.pretty_repr())
    return EventSourceResponse(
        generate(),
        media_type="text/event-stream"
    )

@router.get("/health")
async def health_check():
    return {"status": "healthy"}


app = FastAPI()
app.include_router(router)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)