//! Knowledge Graph Entity Types and Structures
//!
//! Defines entity types, relationships, and properties for the knowledge graph.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Entity types in the knowledge graph
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum EntityType {
    /// Person entity
    Person,
    /// Organization entity
    Organization,
    /// Location entity
    Location,
    /// Concept/Topic entity
    Concept,
    /// Event entity
    Event,
    /// Document entity
    Document,
    /// Generic/Unknown entity
    Generic,
}

impl EntityType {
    /// Get string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            EntityType::Person => "PERSON",
            EntityType::Organization => "ORGANIZATION",
            EntityType::Location => "LOCATION",
            EntityType::Concept => "CONCEPT",
            EntityType::Event => "EVENT",
            EntityType::Document => "DOCUMENT",
            EntityType::Generic => "GENERIC",
        }
    }

    /// Parse from string
    pub fn parse(s: &str) -> Self {
        match s.to_uppercase().as_str() {
            "PERSON" => EntityType::Person,
            "ORGANIZATION" => EntityType::Organization,
            "LOCATION" => EntityType::Location,
            "CONCEPT" => EntityType::Concept,
            "EVENT" => EntityType::Event,
            "DOCUMENT" => EntityType::Document,
            _ => EntityType::Generic,
        }
    }
}

/// Relationship types in the knowledge graph
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum RelationType {
    /// "is related to" relationship
    RelatedTo,
    /// "is part of" relationship
    PartOf,
    /// "contains" relationship
    Contains,
    /// "created by" relationship
    CreatedBy,
    /// "located in" relationship
    LocatedIn,
    /// "works for" relationship
    WorksFor,
    /// "participates in" relationship
    ParticipatesIn,
    /// "mentions" relationship
    Mentions,
    /// "references" relationship
    References,
    /// Custom relationship
    Custom(String),
}

impl RelationType {
    /// Get string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            RelationType::RelatedTo => "RELATED_TO",
            RelationType::PartOf => "PART_OF",
            RelationType::Contains => "CONTAINS",
            RelationType::CreatedBy => "CREATED_BY",
            RelationType::LocatedIn => "LOCATED_IN",
            RelationType::WorksFor => "WORKS_FOR",
            RelationType::ParticipatesIn => "PARTICIPATES_IN",
            RelationType::Mentions => "MENTIONS",
            RelationType::References => "REFERENCES",
            RelationType::Custom(_) => "CUSTOM",
        }
    }

    /// Get custom string if applicable
    pub fn custom_name(&self) -> Option<&str> {
        if let RelationType::Custom(name) = self {
            Some(name)
        } else {
            None
        }
    }
}

/// Entity in the knowledge graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Entity {
    /// Unique identifier
    pub id: String,
    /// Entity name/label
    pub name: String,
    /// Entity type
    pub entity_type: EntityType,
    /// Description
    pub description: Option<String>,
    /// Custom properties
    pub properties: HashMap<String, String>,
    /// Source documents where this entity was found
    pub sources: Vec<String>,
}

impl Entity {
    /// Create a new entity
    pub fn new(id: String, name: String, entity_type: EntityType) -> Self {
        Entity {
            id,
            name,
            entity_type,
            description: None,
            properties: HashMap::new(),
            sources: Vec::new(),
        }
    }

    /// Add description
    pub fn with_description(mut self, desc: String) -> Self {
        self.description = Some(desc);
        self
    }

    /// Add a property
    pub fn with_property(mut self, key: String, value: String) -> Self {
        self.properties.insert(key, value);
        self
    }

    /// Add a source document
    pub fn with_source(mut self, source: String) -> Self {
        self.sources.push(source);
        self
    }

    /// Add multiple sources
    pub fn with_sources(mut self, sources: Vec<String>) -> Self {
        self.sources.extend(sources);
        self
    }
}

/// Relationship/Edge between entities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Relationship {
    /// Source entity ID
    pub source: String,
    /// Target entity ID
    pub target: String,
    /// Relationship type
    pub rel_type: RelationType,
    /// Relationship strength/weight (0.0 - 1.0)
    pub weight: f32,
    /// Optional relationship properties
    pub properties: HashMap<String, String>,
}

impl Relationship {
    /// Create a new relationship
    pub fn new(source: String, target: String, rel_type: RelationType) -> Self {
        Relationship {
            source,
            target,
            rel_type,
            weight: 1.0,
            properties: HashMap::new(),
        }
    }

    /// Set relationship weight
    pub fn with_weight(mut self, weight: f32) -> Self {
        self.weight = weight.clamp(0.0, 1.0);
        self
    }

    /// Add a property
    pub fn with_property(mut self, key: String, value: String) -> Self {
        self.properties.insert(key, value);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_entity_type_as_str() {
        assert_eq!(EntityType::Person.as_str(), "PERSON");
        assert_eq!(EntityType::Organization.as_str(), "ORGANIZATION");
        assert_eq!(EntityType::Location.as_str(), "LOCATION");
    }

    #[test]
    fn test_entity_type_parse() {
        assert_eq!(EntityType::parse("PERSON"), EntityType::Person);
        assert_eq!(EntityType::parse("person"), EntityType::Person);
        assert_eq!(EntityType::parse("INVALID"), EntityType::Generic);
    }

    #[test]
    fn test_entity_creation() {
        let entity = Entity::new("e1".into(), "Alice".into(), EntityType::Person)
            .with_description("A software engineer".into())
            .with_property("role".into(), "engineer".into());

        assert_eq!(entity.id, "e1");
        assert_eq!(entity.name, "Alice");
        assert_eq!(entity.entity_type, EntityType::Person);
        assert_eq!(entity.description, Some("A software engineer".into()));
        assert_eq!(entity.properties.get("role"), Some(&"engineer".into()));
    }

    #[test]
    fn test_relationship_creation() {
        let rel = Relationship::new("e1".into(), "e2".into(), RelationType::WorksFor)
            .with_weight(0.9)
            .with_property("since".into(), "2020".into());

        assert_eq!(rel.source, "e1");
        assert_eq!(rel.target, "e2");
        assert_eq!(rel.weight, 0.9);
        assert_eq!(rel.properties.get("since"), Some(&"2020".into()));
    }

    #[test]
    fn test_relationship_type_custom() {
        let rel_type = RelationType::Custom("knows_well".into());
        assert_eq!(rel_type.as_str(), "CUSTOM");
        assert_eq!(rel_type.custom_name(), Some("knows_well"));
    }
}
