from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from enum import Enum

class TaskType(str, Enum):
    summarize = "summarize"
    classify  = "classify"
    extract   = "extract"
    translate = "translate"

class CompletionRequest(BaseModel):
    prompt: str = Field(..., min_length=1, max_length=10000)
    task: TaskType = TaskType.summarize
    target_language: str | None = None  # for translate task

# System prompts per task — defined server-side, not by client!
SYSTEM_PROMPTS = {
    TaskType.summarize: "Summarize the given text concisely in 2-3 sentences.",
    TaskType.classify:  "Classify the sentiment as POSITIVE, NEGATIVE, or NEUTRAL.",
    TaskType.extract:   "Extract key facts as a JSON list of {fact, source} objects.",
    TaskType.translate: "Translate the text to {language}. Output only the translation.",
}

@app.post("/complete")
async def complete(body: CompletionRequest):
    system = SYSTEM_PROMPTS[body.task]
    if body.task == TaskType.translate:
        if not body.target_language:
            raise HTTPException(400, "target_language is required for translate")
        system = system.format(language=body.target_language)

    response = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": body.prompt},
        ]
    )
    return {"result": response.choices[0].message.content}

from fastapi import FastAPI
from pydantic import BaseModel
import openai, asyncio

class EmbedRequest(BaseModel):
    texts: list[str]             # send one or many at once
    model: str = "text-embedding-3-small"

class EmbedResponse(BaseModel):
    embeddings: list[list[float]]  # one vector per input text
    dimensions: int
    model: str

@app.post("/embed", response_model=EmbedResponse)
async def embed(body: EmbedRequest):
    # OpenAI accepts up to 2048 texts in one call
    response = await client.embeddings.create(
        model=body.model,
        input=body.texts,
    )
    vectors = [item.embedding for item in response.data]
    return EmbedResponse(
        embeddings=vectors,
        dimensions=len(vectors[0]),
        model=response.model,
    )

# Usage — check similarity between two texts
import numpy as np

def cosine_similarity(a: list[float], b: list[float]) -> float:
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import asyncio

async def sse_event(data: str, event: str = None) -> str:
    """Format a Server-Sent Event string."""
    lines = []
    if event:
        lines.append(f"event: {event}")
    lines.append(f"data: {data}")
    lines.append("")  # blank line terminates the event
    return "\n".join(lines) + "\n"

@app.get("/chat/sse")
async def chat_sse(prompt: str, request: Request):
    async def generate():
        stream = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            stream=True,
        )
        async for chunk in stream:
            # Check if client disconnected
            if await request.is_disconnected():
                break
            delta = chunk.choices[0].delta.content
            if delta:
                yield await sse_event(
                    data=json.dumps({"token": delta}),
                    event="token"
                )
        yield await sse_event(data="[DONE]", event="done")

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",  # disable nginx buffering!
        }
    )

# Client: const source = new EventSource("/chat/sse?prompt=Hello");
# source.addEventListener("token", e => appendText(JSON.parse(e.data).token));
# source.addEventListener("done", () => source.close());

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import json

@app.websocket("/ws/chat")
async def ws_chat(websocket: WebSocket):
    await websocket.accept()
    # Keep track of conversation history
    history: list[dict] = []

    try:
        while True:
            # 1. Receive user message
            data = await websocket.receive_text()
            user_msg = json.loads(data)["message"]
            history.append({"role": "user", "content": user_msg})

            # 2. Stream LLM response back token by token
            full_response = ""
            stream = await client.chat.completions.create(
                model="gpt-4o-mini",
                messages=history,
                stream=True
            )
            async for chunk in stream:
                delta = chunk.choices[0].delta.content
                if delta:
                    full_response += delta
                    await websocket.send_json({
                        "type": "token",
                        "content": delta
                    })

            # 3. Signal end; add assistant reply to history
            await websocket.send_json({"type": "done"})
            history.append({"role": "assistant", "content": full_response})

    except WebSocketDisconnect:
        pass  # Client disconnected cleanly

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from enum import Enum
import uuid, redis.asyncio as redis, json

class AgentStatus(str, Enum):
    idle    = "idle"
    running = "running"
    waiting = "waiting"  # waiting for human input
    done    = "done"
    error   = "error"

class AgentSession(BaseModel):
    session_id: str
    status: AgentStatus
    messages: list[dict]  # full history
    current_task: str | None = None

redis_client = redis.Redis(host="localhost", decode_responses=True)

# Create a new agent session
@app.post("/agent/sessions")
async def create_session(initial_prompt: str):
    session_id = str(uuid.uuid4())
    session = AgentSession(
        session_id=session_id,
        status=AgentStatus.idle,
        messages=[{"role": "user", "content": initial_prompt}]
    )
    await redis_client.setex(
        f"session:{session_id}", 3600,  # 1hr TTL
        session.model_dump_json()
    )
    return {"session_id": session_id}

# Get current agent state
@app.get("/agent/sessions/{session_id}")
async def get_session(session_id: str):
    data = await redis_client.get(f"session:{session_id}")
    if not data:
        raise HTTPException(404, "Session not found or expired")
    return AgentSession.model_validate_json(data)

# Send a message to an existing session
@app.post("/agent/sessions/{session_id}/messages")
async def send_message(session_id: str, message: str):
    session = await get_session(session_id)
    session.messages.append({"role": "user", "content": message})
    session.status = AgentStatus.running
    # Persist and dispatch to background task ...
    return {"status": "queued"}

from fastapi import FastAPI
from pydantic import BaseModel
import asyncio

class BlogRequest(BaseModel):
    topic: str
    target_audience: str
    word_count: int = 800

class BlogResponse(BaseModel):
    research:  str  # Step 1 output
    outline:   str  # Step 2 output
    draft:     str  # Step 3 output
    final:     str  # Step 4 output

async def llm(system: str, prompt: str) -> str:
    """Thin helper to call LLM with a system + user message."""
    r = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": prompt},
        ]
    )
    return r.choices[0].message.content

@app.post("/workflow/blog", response_model=BlogResponse)
async def blog_workflow(body: BlogRequest):
    # Step 1: Research key points
    research = await llm(
        "You are a research assistant. List 5 key facts about the topic.",
        f"Topic: {body.topic} | Audience: {body.target_audience}"
    )

    # Step 2: Create outline using research
    outline = await llm(
        "You are a content strategist. Create a blog outline.",
        f"Topic: {body.topic}\nResearch:\n{research}"
    )

    # Step 3: Write draft based on outline
    draft = await llm(
        f"Write a {body.word_count}-word blog post for {body.target_audience}.",
        f"Outline:\n{outline}\nResearch:\n{research}"
    )

    # Step 4: Polish the draft
    final = await llm(
        "Polish this blog post: fix grammar, improve flow, add a strong CTA.",
        draft
    )

    return BlogResponse(
        research=research, outline=outline, draft=draft, final=final
    )

from dataclasses import dataclass

@dataclass
class Chunk:
    doc_id:   str
    chunk_id: str
    text:     str
    start:    int   # character position in source
    metadata: dict

def split_into_chunks(
    text: str,
    doc_id: str,
    chunk_size: int = 500,   # tokens (~2000 chars)
    overlap:    int = 50,    # overlap to preserve context
) -> list[Chunk]:
    words   = text.split()
    chunks  = []
    i       = 0

    while i < len(words):
        window      = words[i : i + chunk_size]
        chunk_text  = " ".join(window)
        chunk_start = len(" ".join(words[:i]))

        chunks.append(Chunk(
            doc_id   = doc_id,
            chunk_id = f"{doc_id}_{i}",
            text     = chunk_text,
            start    = chunk_start,
            metadata = {"word_index": i, "word_count": len(window)}
        ))
        i += chunk_size - overlap  # slide window with overlap

    return chunks

async def index_document(doc_id: str, filename: str, text: str):
    # 1. Split into chunks
    chunks = split_into_chunks(text, doc_id)

    # 2. Embed ALL chunks in one API call (batch)
    texts_to_embed = [c.text for c in chunks]
    embed_response = await client.embeddings.create(
        model="text-embedding-3-small",
        input=texts_to_embed
    )
    vectors = [item.embedding for item in embed_response.data]

    # 3. Store in vector DB (pseudocode — depends on your DB choice)
    for chunk, vector in zip(chunks, vectors):
        await vector_db.upsert(
            id=chunk.chunk_id,
            values=vector,
            metadata={"text": chunk.text, "doc_id": doc_id, "filename": filename}
        )

from fastapi import FastAPI
from pydantic import BaseModel

class SearchResult(BaseModel):
    chunk_id:   str
    text:       str
    score:      float   # cosine similarity (0-1, higher = more similar)
    doc_id:     str
    filename:   str

class SearchResponse(BaseModel):
    query:   str
    results: list[SearchResult]

@app.get("/search", response_model=SearchResponse)
async def search(q: str, top_k: int = 5, doc_id: str | None = None):
    # 1. Embed the search query
    embed_response = await client.embeddings.create(
        model="text-embedding-3-small",
        input=[q]
    )
    query_vector = embed_response.data[0].embedding

    # 2. Search vector DB for nearest neighbors
    filter = {"doc_id": doc_id} if doc_id else {}
    matches = await vector_db.query(
        vector=query_vector,
        top_k=top_k,
        include_metadata=True,
        filter=filter
    )

    # 3. Format results
    results = [
        SearchResult(
            chunk_id=m.id,
            text=m.metadata["text"],
            score=m.score,
            doc_id=m.metadata["doc_id"],
            filename=m.metadata["filename"]
        )
        for m in matches
    ]
    return SearchResponse(query=q, results=results)

from fastapi import FastAPI
from pydantic import BaseModel

class RAGRequest(BaseModel):
    question: str
    doc_id:   str | None = None  # restrict to one document
    top_k:    int = 5

class RAGResponse(BaseModel):
    answer:  str
    sources: list[dict]  # chunks used to generate the answer

@app.post("/rag/chat", response_model=RAGResponse)
async def rag_chat(body: RAGRequest):
    # ── Step 1: Retrieve relevant chunks ─────────────────────
    embed_r = await client.embeddings.create(
        model="text-embedding-3-small", input=[body.question]
    )
    query_vec = embed_r.data[0].embedding

    matches = await vector_db.query(
        vector=query_vec, top_k=body.top_k,
        include_metadata=True,
        filter={"doc_id": body.doc_id} if body.doc_id else {}
    )

    # ── Step 2: Build context from retrieved chunks ───────────
    context_parts = []
    sources       = []

    for i, match in enumerate(matches, 1):
        context_parts.append(f"[Source {i}] {match.metadata['text']}")
        sources.append({
            "rank":     i,
            "filename": match.metadata["filename"],
            "score":    round(match.score, 3),
            "excerpt":  match.metadata["text"][:200] + "..."
        })

    context = "\n\n".join(context_parts)

    # ── Step 3: Call LLM with context in system prompt ───────
    system = f"""You are a helpful assistant. Answer the user's question
ONLY using the provided context below. If the answer is not in the context,
say "I don't have enough information in the provided documents."

CONTEXT:
{context}"""

    response = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": body.question},
        ]
    )

    return RAGResponse(
        answer=response.choices[0].message.content,
        sources=sources
    )