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
#from prompts import PRESENTATION_SYSTEM_PROMPT

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]:
    """Tool to create a full slide with the proposed content plan. Create content according to the plan, include images <!-- Title Slide -->
    <style scoped>
    </style>
    # Title
    <div class="subtitle"></div>
    <div class="logo"></div>
    <div class="date"></div>
    ![bg right:40%](https://picsum.photos/400/600)
    ---
    add styles to design a beutiful presentation, add a modern feel, add images using marp syntax and https://picsum.photos/200/300 (width & height)
    """
    # 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. Wait for user input to proceed with the plan or update the current plan based on input.
    3. Then use create_slide tool for each slide in sequence
    4. 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"]
        })
        print(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)