from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import psycopg2
import httpx
import json
import os

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://apply4jobs.de", "https://www.apply4jobs.de", "http://localhost:3000"],
    allow_methods=["POST", "GET", "OPTIONS"],
    allow_headers=["*"],
)

model = SentenceTransformer("intfloat/multilingual-e5-small")

DB_CONFIG = {
    "host": "postgres",
    "port": 5432,
    "dbname": "anythingllm",
    "user": "anythingllm",
    "password": "any0203thing78llm"
}

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
OPENAI_MODEL = "gpt-4o-mini"

SYSTEM_PROMPT = """Du bist ein professioneller Assistent für Sebastian Fröhlich.
Beantworte ausschließlich Fragen zu seiner Person, seinen Fähigkeiten,
Projekten und Berufserfahrung. Nutze nur die bereitgestellten Dokumente.
Antworte auf Deutsch oder Englisch je nach Sprache des Recruiters.
Wenn du eine Frage nicht aus den Dokumenten beantworten kannst, sage das ehrlich.
Antworte immer in vollständigen, professionellen Sätzen."""

def get_context(query: str, top_k: int = 5) -> str:
    query_text = f"query: {query}"
    embedding = model.encode(query_text).tolist()

    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    cur.execute("""
        SELECT metadata->>'text'
        FROM anythingllm_vectors
        WHERE namespace = 'mein-workspace'
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """, (embedding, top_k))
    rows = cur.fetchall()
    cur.close()
    conn.close()

    return "\n\n---\n\n".join([row[0] for row in rows])


class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: str = OPENAI_MODEL
    messages: list[Message]
    stream: bool = False

class QueryRequest(BaseModel):
    query: str
    top_k: int = 5


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/v1/models")
def list_models():
    return {
        "object": "list",
        "data": [{
            "id": "sebastian-rag",
            "object": "model",
            "created": 1700000000,
            "owned_by": "apply4jobs"
        }]
    }


@app.post("/retrieve")
def retrieve(req: QueryRequest):
    context = get_context(req.query, req.top_k)
    return {"context": context, "system_prompt": SYSTEM_PROMPT}


@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
    user_message = next(
        (m.content for m in reversed(req.messages) if m.role == "user"), ""
    )

    context = get_context(user_message)

    enriched_messages = [
        {
            "role": "system",
            "content": f"{SYSTEM_PROMPT}\n\n## Relevante Dokumente:\n{context}"
        },
        *[{"role": m.role, "content": m.content} for m in req.messages]
    ]

    async def stream_response():
        async with httpx.AsyncClient(timeout=60) as client:
            async with client.stream(
                "POST",
                "https://api.openai.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {OPENAI_API_KEY}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": OPENAI_MODEL,
                    "messages": enriched_messages,
                    "stream": True
                }
            ) as response:
                async for chunk in response.aiter_bytes():
                    yield chunk

    if req.stream:
        return StreamingResponse(
            stream_response(),
            media_type="text/event-stream"
        )
    else:
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                "https://api.openai.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {OPENAI_API_KEY}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": OPENAI_MODEL,
                    "messages": enriched_messages,
                    "stream": False
                }
            )
            return response.json()