"""Pydantic schemas for res_sum — entities, relationships, summaries, and domain config."""

from __future__ import annotations

from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field


# ---------------------------------------------------------------------------
# Knowledge Graph Node Types (research-level, not keyword-level)
# ---------------------------------------------------------------------------

class NodeType(str, Enum):
    """Node types for the per-paper knowledge graph."""
    RESEARCH_PROBLEM = "RESEARCH_PROBLEM"   # The question the paper addresses
    FINDING = "FINDING"                      # An empirical result with evidence
    CLAIM = "CLAIM"                          # An interpretive assertion by the authors
    METHOD = "METHOD"                        # The approach used (specific, not generic)
    STUDY_SYSTEM = "STUDY_SYSTEM"            # Species + location + context
    LIMITATION = "LIMITATION"                # Acknowledged caveats
    PAPER = "PAPER"                          # The paper itself as a node


class EdgeType(str, Enum):
    """Edge types for within-paper relationships."""
    HAS_FINDING = "HAS_FINDING"              # Paper → Finding
    HAS_CLAIM = "HAS_CLAIM"                  # Paper → Claim
    ADDRESSES = "ADDRESSES"                   # Paper → Research Problem
    USES_METHOD = "USES_METHOD"               # Paper → Method
    STUDIED_IN = "STUDIED_IN"                 # Paper → Study System
    SUPPORTED_BY = "SUPPORTED_BY"             # Claim → Finding (this claim is backed by this finding)
    PRODUCED_BY = "PRODUCED_BY"               # Finding → Method (this method produced this finding)
    HAS_LIMITATION = "HAS_LIMITATION"         # Paper → Limitation


# ---------------------------------------------------------------------------
# Extracted content schemas (what the LLM produces per paper)
# ---------------------------------------------------------------------------

class ExtractedEntity(BaseModel):
    """A knowledge node extracted from a paper."""
    name: str = Field(..., description="Short descriptive label for this node")
    node_type: str = Field(..., description="Node type (RESEARCH_PROBLEM, FINDING, CLAIM, METHOD, STUDY_SYSTEM, LIMITATION)")
    description: str = Field(default="", description="Full description of this node — the actual content")
    evidence: str = Field(default="", description="Supporting statistical evidence, if applicable")
    confidence: str = Field(default="", description="Evidence strength: strong, moderate, weak")
    source_chunk_id: str = Field(default="", description="ID of the source chunk in ChromaDB")
    # Keep aliases for backward compat
    aliases: list[str] = Field(default_factory=list, description="Alternative phrasings")
    # Legacy field mapping
    entity_type: str = Field(default="", description="Deprecated — use node_type instead")


class ExtractedRelationship(BaseModel):
    """A relationship between two knowledge nodes within a paper."""
    source: str = Field(..., description="Source node name")
    target: str = Field(..., description="Target node name")
    relation_type: str = Field(..., description="Relationship type (HAS_FINDING, SUPPORTED_BY, etc.)")
    evidence: str = Field(default="", description="Brief explanation of why this relationship holds")
    source_chunk_id: str = Field(default="", description="ID of the source chunk in ChromaDB")


class ExtractionResult(BaseModel):
    """Combined extraction output from a paper (or chunk of a paper)."""
    entities: list[ExtractedEntity] = Field(default_factory=list)
    relationships: list[ExtractedRelationship] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# Domain config schemas
# ---------------------------------------------------------------------------

class EntityTypeDef(BaseModel):
    """Definition of a node type within a domain config."""
    name: str = Field(..., description="Node type name (e.g. FINDING)")
    description: str = Field(..., description="What this node type represents")
    examples: list[str] = Field(default_factory=list, description="Example nodes")


class DomainConfig(BaseModel):
    """Configuration for a scientific domain — defines node types, relationships, and prompts."""
    name: str = Field(..., description="Domain name (e.g. ecology, general)")
    entity_types: list[EntityTypeDef] = Field(default_factory=list)
    relationship_types: list[str] = Field(default_factory=list)
    cot_prompt: str = Field(default="", description="Chain-of-Thought prompt template for this domain")
    section_headers: list[str] = Field(
        default_factory=list,
        description="Additional section headers to recognize in PDFs",
    )


# ---------------------------------------------------------------------------
# Summary output schemas
# ---------------------------------------------------------------------------

class StructuredSummary(BaseModel):
    """Structured summary output from a single paper."""
    paper_filename: str = Field(default="")
    research_question: str = Field(default="")
    study_species: list[str] = Field(default_factory=list)
    study_location: str = Field(default="")
    temporal_scale: str = Field(default="")
    spatial_scale: str = Field(default="")
    methods: str = Field(default="")
    key_findings: list[str] = Field(default_factory=list)
    statistical_evidence: list[str] = Field(default_factory=list)
    ecological_implications: str = Field(default="")
    limitations: str = Field(default="")
    conservation_relevance: str = Field(default="")
    raw_summary: str = Field(default="", description="Free-text summary if structured parsing fails")


class PaperMetadata(BaseModel):
    """Metadata about an ingested paper."""
    filename: str
    filepath: str = ""
    num_chunks: int = 0
    num_entities: int = 0
    num_relationships: int = 0
    sections_found: list[str] = Field(default_factory=list)
    ingested: bool = False
