from dataclasses import dataclass, field
from typing import List


@dataclass
class Incident:
    """Represents an incident record in the VDB collection."""

    id: str
    vector: List[float]
    company: str
    year: int
    narrative: str = ""
    impact: str = ""
    industries: List[str] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> "Incident":
        """Create an Incident from a dictionary, handling Milvus/Zilliz types."""
        # Merge trigger_keywords into tags for backward compatibility
        tags = list(data.get("tags", []))
        trigger_keywords = list(data.get("trigger_keywords", []))
        merged_tags = list(set(tags + trigger_keywords))

        return cls(
            id=data.get("id"),
            vector=list(data.get("vector", [])),
            company=data.get("company", ""),
            year=data.get("year", 0),
            narrative=data.get("narrative", ""),
            impact=data.get("impact", ""),
            industries=list(data.get("industries", [])),
            tags=merged_tags,
        )

    def to_dict(self) -> dict:
        """Convert to dictionary for insertion."""
        return {
            "id": self.id,
            "vector": self.vector,
            "company": self.company,
            "year": self.year,
            "narrative": self.narrative,
            "impact": self.impact,
            "industries": self.industries,
            "tags": self.tags,
        }


@dataclass
class IncidentQueryResult:
    """Represents a lightweight incident record for search/query results."""

    id: str
    company: str
    year: int
    narrative: str = ""
    impact: str = ""
    industries: List[str] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> "IncidentQueryResult":
        """Create a QueryResult from a dictionary, handling Milvus/Zilliz types."""
        # Merge trigger_keywords into tags for backward compatibility
        tags = list(data.get("tags", []))
        trigger_keywords = list(data.get("trigger_keywords", []))
        merged_tags = list(set(tags + trigger_keywords))

        return cls(
            id=data.get("id"),
            company=data.get("company", ""),
            year=data.get("year", 0),
            narrative=data.get("narrative", ""),
            impact=data.get("impact", ""),
            industries=list(data.get("industries", [])),
            tags=merged_tags,
        )

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "id": self.id,
            "company": self.company,
            "year": self.year,
            "narrative": self.narrative,
            "impact": self.impact,
            "industries": self.industries,
            "tags": self.tags,
        }

    def to_prompt_str(self) -> str:
        """Returns a string representation of the incident query result."""
        return (
            f"Incident: {self.company} ({self.year})\n"
            f"Narrative: {self.narrative}\n"
            f"Impact: {self.impact}\n"
            f"Tags: {', '.join(self.tags)}\n"
            f"Industries: {', '.join(self.industries)}\n"
        )
