from enum import Enum

import ollama
from pydantic import BaseModel

class Category(str, Enum):
    REGULATORY_NEWS = "Regulatory News"
    INSTITUTIONAL_ADOPTION = "Institutional Adoption"
    MARKET_SENTIMENT = "Market Sentiment"
    MACROECONOMIC_FACTORS = "Macroeconomic Factors"
    SECURITY_HACKS = "Security & Hacks"
    TECHNOLOGICAL_DEVELOPMENTS = "Technological Developments"
    WHALE_EXCHANGE_ACTIVITY = "Whale & Exchange Activity"

class ArticleClassification(BaseModel):
    category: Category
    sentiment: int

class ArticleAnalyzer:
    def __init__(self):
        self.base_prompt = """
            Classify the following article into one of these categories: 
            - Regulatory News
            - Institutional Adoption
            - Market Sentiment
            - Macroeconomic Factors
            - Security & Hacks
            - Technological Developments
            - Whale & Exchange Activity

            Also, assign a sentiment (1 for Positive, -1 for Negative, or 0 for Neutral).
        """
        print(f"This JSON model is going to be used for structured ouput classification : {ArticleClassification.model_json_schema()}")

    def classify_article(self, article_text):
        prompt = f"""{self.base_prompt}

        ARTICLE: {article_text}

        OUTPUT FORMAT:
        Category: <category>
        Sentiment: <sentiment>
        """
        response = ollama.chat(model="llama3.2",
                               messages=[{"role": "user", "content": prompt}],
                               format=ArticleClassification.model_json_schema())
        return response['message']['content']