import dspy
from typing import List
from config_loader import load_config

CFG = load_config()
INGESTION_CONFIG = CFG["ingestion_agent"]


class IngestionSignature(dspy.Signature):
    f"{INGESTION_CONFIG['ingestion_signature']}"

    note: str = dspy.InputField(desc="The DM notes or session recap content.")
    answer: dict[str, str | List] = dspy.OutputField(
        desc="the metadata dictionary with the keys; synopsis, tags, entities, relationships"
    )


class IngestionAgent(dspy.Module):
    def __init__(self):
        self.ingest = dspy.Predict(IngestionSignature)

    def ingest_with_relationships(self, note: str) -> dict:
        """Ingest notes and return metadata including extracted relationships."""
        response = self.ingest(note=note)
        result = response.answer

        if not isinstance(result, dict):
            result = {
                "synopsis": "Failed to parse",
                "tags": [],
                "entities": [],
                "relationships": [],
            }

        if "relationships" not in result:
            entities = result.get("entities", [])
            relationships = []

            for i, ent1 in enumerate(entities):
                for ent2 in entities[i + 1 :]:
                    relationships.append(
                        {"entity1": ent1, "entity2": ent2, "type": "co-occurs_with", "strength": 1}
                    )

            result["relationships"] = relationships

        return result