from fastapi import FastAPI
from pydantic import BaseModel
import spacy
from collections import Counter
from fastapi.middleware.cors import CORSMiddleware

nlp = spacy.load("en_core_web_sm")
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

class TextPayload(BaseModel):
    text: str

@app.post("/analyze")
async def analyze_text(payload: TextPayload):
    doc = nlp(payload.text)
    entity_freq = Counter([ent.text.lower() for ent in doc.ents])
    
    brand_safety = ['johnson & johnson', 'at&t', 'h&m']
    trash_words = ['zero', 'one', 'two', 'three', 'privacy & cookie', 'accept decline', 'ip']

    processed_entities = []
    seen = set()

    for ent in doc.ents:
        raw_text = ent.text.strip()
        e_lower = raw_text.lower()
        
        # 1. Filter out trash, numbers, and very short 2-letter acronyms from being ORGs
        if e_lower in trash_words or raw_text.isdigit():
            continue

        sub_items = [raw_text]
        if e_lower not in brand_safety:
            if " and " in e_lower or " & " in raw_text:
                parts = raw_text.replace(" and ", " & ").split(" & ")
                if all(1 <= len(p.strip().split()) <= 2 for p in parts):
                    sub_items = [p.strip() for p in parts if len(p.strip()) > 1]

        for final_text in sub_items:
            f_lower = final_text.lower()
            if f_lower in seen or len(final_text) < 2:
                continue

            # Repetition filter
            words = f_lower.split()
            if len(words) > 1 and words[-1] == words[0]:
                continue

            label = ent.label_
            
            # Force Category logic
            concepts = ['privacy', 'cookies', 'cookie', 'booster', 'get', 'noticed', 'free', 'social media', 'ip']
            tech_items = ['exif', 'mediaif', 'cookiesif', 'json', 'api', 'schema', 'vps', 'server', 'gravatar']

            if any(c == f_lower for c in concepts): # Direct match for short words like IP
                label = "CONCEPT"
            elif any(t in f_lower for t in tech_items):
                label = "PRODUCT"
            
            # Final check: 2-letter words like 'IP' should almost never be an 'ORG'
            if len(final_text) <= 2 and label == "ORG":
                label = "CONCEPT"

            if len(final_text.split()) >= 4:
                continue

            processed_entities.append({
                "text": final_text,
                "label": label,
                "salience_score": round(entity_freq.get(f_lower, 1) / len(doc.ents), 4) if len(doc.ents) > 0 else 0,
                "frequency": entity_freq.get(f_lower, 1)
            })
            seen.add(f_lower)

    sorted_entities = sorted(processed_entities, key=lambda x: x['frequency'], reverse=True)
    return {"status": "success", "entities": sorted_entities[:50]}