from typing import List, Dict, Any, Optional, Annotated
from pydantic import BaseModel, Field
from langchain_core.messages import BaseMessage
from langgraph.graph import add_messages
class ChatbotState(BaseModel):
    def get(self, key, default=None):
        """
        Allow dict-like .get() access for compatibility.
        """
        # First try attribute directly
        if hasattr(self, key):
            return getattr(self, key)
        # Fallback: check if it's in __dict__
        return self.__dict__.get(key, default)
    def setdefault(self, key, default):
        """
        Dict-like setdefault: if attribute is None, set it to default and return it.
        Otherwise, return existing value.
        """
        if hasattr(self, key):
            value = getattr(self, key)
            if value is None:
                setattr(self, key, default)
                return default
            return value
        else:
            # attribute does not exist: set it
            setattr(self, key, default)
            return default
    profile: Dict[str, Any] = Field(..., description="Preprocessed / summarized profile data")
    profile_url: Optional[str] = Field(
        default=None,
        description="Original LinkedIn profile URL provided by the user."
    )

    # Quick access sections (about, headline, skills etc.)
    sections: Dict[str, str] = Field(..., description="Flattened profile sections for quick access")

    # Enhancements and analysis results
    enhanced_content: Dict[str, str] = Field(
    default_factory=dict,
    description=(
        "Map of improved or rewritten profile sections generated by the ContentGenerator tool. "
        "Keys are section names (e.g., 'about', 'headline'); values are enhanced text."
    )
)

    profile_analysis: Optional[Dict[str, Any]] = Field(
        None,
        description=(
            "Structured analysis of the user's profile produced by the ProfileAnalyzer tool, "
            "including strengths, weaknesses, and actionable suggestions."
        )
    )

    job_fit: Optional[Dict[str, Any]] = Field(
        None,
        description=(
            "Assessment result from the JobMatcher tool, detailing how well the user's profile matches "
            "the target role, including missing skills and match score."
        )
    )

    target_role: Optional[str] = Field(
        None,
        description=(
            "Target job role the user is aiming for. "
            "Can be set by the user directly during the conversation or inferred by the chatbot."
        )
    )

    editing_section: Optional[str] = Field(
        None,
        description=(
            "Name of the profile section currently being edited or improved, "
            "set dynamically when the ContentGenerator tool is invoked."
        )
    )
    next_tool_name: Optional[str] = Field(
        default=None,
        description="Name of the next tool the chatbot wants to call, set dynamically after LLM response."
    )


    # Annotated chat history directly using BaseMessage
    messages: Annotated[List[BaseMessage], add_messages] = Field(
    default_factory=list,
    description="List of user and assistant messages"
)



class ProfileAnalysisStrengths(BaseModel):
    technical: List[str]
    projects: List[str]
    education: List[str]
    soft_skills: List[str]

class ProfileAnalysisWeaknesses(BaseModel):
    technical_gaps: List[str]
    project_or_experience_gaps: List[str]
    missing_context: List[str]

class ProfileAnalysisModel(BaseModel):
    strengths: ProfileAnalysisStrengths
    weaknesses: ProfileAnalysisWeaknesses
    suggestions: List[str]

class JobFitModel(BaseModel):
    match_score: int = Field(..., ge=0, le=100)
    missing_skills: List[str]
    suggestions: List[str]

class ContentGenerationModel(BaseModel):
    new_content: str


# ========== 6. MEMORY SETUP ==========

class UserMemory:
    def __init__(self):
        self.profile = None
        self.target_roles = []
        self.history = []

    def save(self, key, value):
        self.history.append((key, value))

    def get_history(self):
        return self.history