import json
from app.config import settings

CATEGORIES = ["Politics","National","International","Sports","Business","Technology",
    "Education","Health","Science","Entertainment","Crime","Court","Weather",
    "Editorial","Opinion","Local News","State News","Environment","Lifestyle",
    "Automobile","Travel","Agriculture","Jobs","Classifieds","Others"]

SUBCATS = {"Sports": ["Cricket","Football","Tennis","Others"],
           "Business": ["Finance","Markets","Economy"],
           "Technology": ["AI","Gadgets","Software"],
           "Entertainment": ["Movies","TV","Music"],
           "Politics": ["State Politics","National Politics","International"]}

PROMPT = """You are a news classifier. Given the headline and text, return strict JSON:
{{"category": one of {cats}, "subcategory": string, "confidence": 0-1,
"summary": 2-sentence summary, "keywords": [up to 8],
"persons": [], "locations": [], "organizations": [], "dates": [], "sentiment": "pos|neu|neg"}}

{few_shot_context}
HEADLINE: {headline}
TEXT: {text}
"""

class Categorizer:
    def __init__(self):
        self.provider = settings.LLM_PROVIDER
        if self.provider == "gemini":
            import google.generativeai as genai
            genai.configure(api_key=settings.LLM_API_KEY)
            self.model = genai.GenerativeModel("gemini-1.5-flash")

    def classify(self, headline, text, few_shot_examples=None):
        few_shot_context = ""
        if few_shot_examples:
            few_shot_context = "Here are some past examples of how similar articles were classified by a human:\n\n"
            for i, ex in enumerate(few_shot_examples):
                few_shot_context += f"EXAMPLE {i+1}:\nHEADLINE: {ex['headline']}\nTEXT: {ex['text'][:500]}...\nCLASSIFICATION (JSON): {json.dumps(ex['classification'])}\n\n"
            few_shot_context += "Now, classify the following new article based on these patterns:\n\n"

        prompt = PROMPT.format(cats=CATEGORIES, few_shot_context=few_shot_context, headline=headline, text=text[:4000])
        try:
            if self.provider == "gemini":
                resp = self.model.generate_content(prompt)
                raw = resp.text.strip().strip("```json").strip("```")
                return json.loads(raw)
        except Exception as e:
            return self._fallback(headline, text)

    def _fallback(self, headline, text):
        # Zero-shot via local model if LLM unavailable
        from transformers import pipeline
        clf = pipeline("zero-shot-classification",
                       model="facebook/bart-large-mnli")
        res = clf(headline + ". " + text[:500], CATEGORIES)
        return {"category": res["labels"][0], "subcategory": "",
                "confidence": res["scores"][0], "summary": "", "keywords": [],
                "persons": [], "locations": [], "organizations": [],
                "dates": [], "sentiment": "neu"}