from langchain_community.chat_models import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from typing import List, Dict
import os
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"
DEFAULT_MODEL = "meta-llama/llama-3.1-70b-instruct:free"

class ChatManager:
    def __init__(self, temperature: float = 0.7):
        self.chat = ChatOpenAI(
            openai_api_base=OPENROUTER_API_BASE,
            openai_api_key=OPENROUTER_API_KEY,
            model_name=DEFAULT_MODEL,
            temperature=temperature,
        )
        
        self.system_message = """Du er en avanceret assistent designet til at hjælpe brugere ved at hente den mest relevante information fra et foruddefineret sæt af dokumenter eller sager og derefter give et præcist svar baseret på disse data.

Din opgave er som følger:

1. Når brugeren indsender en forespørgsel, skal du matche forespørgslen med den mest relevante sag fra databasen.
2. Uddrag kun den støjfri og kontekstuelle relevante tekst fra den pågældende sag.
3. Brug den udtrukne tekst til præcist og klart at besvare brugerens forespørgsel.
4. Hvis den relevante tekst ikke findes, informer brugeren om, at informationen ikke er tilgængelig, eller bed om yderligere afklaring.
5. Undgå at give information uden for den udtrukne teksts rækkevidde.

Fokuser altid på relevans og klarhed i dit svar, og oprethold sammenhæng med brugerens oprindelige forespørgsel."""
        self.user_message = """Kontekst: {context}

Forespørgsel: {query}
"""
        
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", self.system_message),
            ("human", self.user_message),
        ])

    def get_response(self, context: str, query: str) -> str:
        prompt_value = self.prompt.invoke({
            "context": context,
            "query": query
        })
        messages = prompt_value.to_messages()
        response = self.chat(messages)
        return response.content