import requests
import json


class API:
    @staticmethod
    def process_text(prompt, model):
        ollama_url = "http://localhost:11434"

        response = requests.post(
            f"{ollama_url}/api/generate", json={"model": model, "prompt": prompt}
        )

        if response.status_code == 200:
            response_data = []
            for line in response.iter_lines():
                line_decoded = line.decode("utf-8")
                line_raw = json.loads(line_decoded)
                response_data.append(line_raw["response"])

            final_response = "".join(response_data)
            generated_text = final_response.splitlines()[-1]
            return generated_text
        else:
            return "Error: " + response.text

    def send_message(self, message, model):
        if model == 1:
            answer = self.process_text(message, "phi3.5")
        elif model == 2:
            answer = self.process_text(message, "gemma2:2b")
        elif model == 3:
            answer = self.process_text(message, "qwen2:0.5b")
        elif model == 4:
            answer = self.process_text(message, "codegemma:2b")
        else:
            return "Invalid choice"
        return answer