# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
"""Extraction d'un GRAPHE DE CONNAISSANCES orienté CRÉATEURS / INFLUENCEURS (ka6).

Objectif initial du flagship : cartographier l'écosystème des créateurs de contenu
et influenceurs québécois — leurs plateformes, handles, nombre d'abonnés, niche,
langues — et leurs relations (agence qui les représente, marques avec qui ils
collaborent, plateformes où ils publient, collectifs dont ils sont membres).

On extrait aussi les entités classiques (entreprises, organisations, agences).
Données PUBLIQUES uniquement ; aucune donnée sensible ou privée.
"""

from __future__ import annotations

from typing import Any

from anthropic import Anthropic

GRAPH_TOOL = {
    "name": "record_graph",
    "description": "Enregistre les entités (dont créateurs/influenceurs) et leurs relations.",
    "input_schema": {
        "type": "object",
        "properties": {
            "entities": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "temp_id": {"type": "string", "description": "Identifiant local unique (e1, e2…) référencé par les relations"},
                        "type": {"type": "string",
                                 "enum": ["creator", "business", "organization", "agency", "brand", "person", "website"],
                                 "description": "Utilise 'creator' pour un influenceur/créateur de contenu, 'agency' pour une agence d'influence/talent, 'brand' pour une marque."},
                        "name": {"type": "string"},
                        "canonical_name": {"type": "string"},
                        "handle": {"type": "string", "description": "Pseudo principal (ex. @nomcreateur)"},
                        "platform": {"type": "string", "description": "Plateforme principale (Instagram, TikTok, YouTube, Twitch, X, Facebook, Snapchat…)"},
                        "followers": {"type": "string", "description": "Nb d'abonnés principal (ex. '125K', '1.2M', 34000)"},
                        "niche": {"type": "string", "description": "Créneau/thématique (mode, beauté, gaming, cuisine, humour, lifestyle, sport, tech, voyage, famille, musique…)"},
                        "languages": {"type": "array", "items": {"type": "string"}, "description": "Langues du contenu (fr, en…)"},
                        "sector": {"type": "string", "description": "Secteur pour les entreprises/agences/marques"},
                        "description": {"type": "string"},
                        "website": {"type": "string"},
                        "email": {"type": "string"},
                        "phone": {"type": "string"},
                        "city": {"type": "string"},
                        "region": {"type": "string", "description": "Ville/région du Québec"},
                        "social_links": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "platform": {"type": "string"},
                                    "url": {"type": "string"},
                                    "followers": {"type": "string", "description": "Abonnés sur CETTE plateforme si connu"},
                                },
                                "required": ["platform", "url"],
                            },
                        },
                        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
                    },
                    "required": ["temp_id", "type", "name"],
                },
            },
            "relations": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "from": {"type": "string"},
                        "to": {"type": "string"},
                        "type": {
                            "type": "string",
                            "enum": ["REPRESENTED_BY", "COLLABORATES_WITH", "CREATES_ON", "SPONSORED_BY",
                                     "PROMOTES", "MEMBER_OF", "MANAGES", "APPEARS_WITH",
                                     "WORKS_AT", "FOUNDER_OF", "PARTNER_OF", "AFFILIATED_WITH"],
                            "description": "Ex : créateur REPRESENTED_BY agence ; créateur COLLABORATES_WITH marque ; créateur CREATES_ON plateforme ; créateur MEMBER_OF collectif.",
                        },
                        "role": {"type": "string", "description": "Précision (ex. 'ambassadrice', 'campagne 2025', 'chaîne principale')"},
                        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
                    },
                    "required": ["from", "to", "type"],
                },
            },
        },
        "required": ["entities"],
    },
}

SYSTEM = (
    "Tu es l'extracteur du bot FLAGSHIP ka6 du Groupe KA. Objectif initial : cartographier "
    "l'écosystème des CRÉATEURS DE CONTENU et INFLUENCEURS québécois à partir de pages publiques "
    "(répertoires, palmarès, rosters d'agences, articles, pages de marques).\n\n"
    "Pour chaque créateur/influenceur (type 'creator') extrais : nom, handle (@), plateforme "
    "principale, nombre d'abonnés, niche/thématique, langues, région, et TOUTES ses plateformes "
    "dans social_links (avec abonnés par plateforme si indiqués).\n"
    "Modélise les RELATIONS : REPRESENTED_BY (agence de talent/influence), COLLABORATES_WITH / "
    "SPONSORED_BY / PROMOTES (marques), CREATES_ON (plateformes), MEMBER_OF (collectif/réseau), "
    "MANAGES (gérant/agence -> créateur).\n"
    "Extrais aussi agences ('agency'), marques ('brand'), entreprises/organisations pertinentes.\n"
    "Règles : données PUBLIQUES uniquement ; priorise le Québec ; n'invente rien (laisse vide si "
    "absent) ; donne un temp_id unique par entité et référence-le dans les relations ; confidence "
    "réaliste. Réponds uniquement via l'outil record_graph."
)


class Extractor:
    def __init__(self, client: Anthropic, model: str):
        self.client = client
        self.model = model

    def extract_graph(self, url: str, title: str, markdown: str) -> dict[str, Any]:
        if not markdown or not markdown.strip():
            return {"entities": [], "relations": []}
        content = markdown[:20000]
        resp = self.client.messages.create(
            model=self.model,
            max_tokens=3500,
            system=SYSTEM,
            tools=[GRAPH_TOOL],
            tool_choice={"type": "tool", "name": "record_graph"},
            messages=[{
                "role": "user",
                "content": f"URL: {url}\nTitre: {title}\n\nContenu (markdown):\n{content}",
            }],
        )
        for block in resp.content:
            if block.type == "tool_use" and block.name == "record_graph":
                data = block.input or {}
                return {"entities": data.get("entities", []) or [],
                        "relations": data.get("relations", []) or []}
        return {"entities": [], "relations": []}