from enum import StrEnum

from pydantic import BaseModel, ConfigDict

from unique_swot.services.collection.schema import Source
from unique_swot.services.generation.models import SWOTExtractionModel


class SWOTComponent(StrEnum):
    STRENGTHS = "Strengths"
    WEAKNESSES = "Weaknesses"
    OPPORTUNITIES = "Opportunities"
    THREATS = "Threats"


class ReportGenerationContext(BaseModel):
    """
    Context information for generating SWOT analysis reports.

    Contains all the necessary information to generate a report for a specific
    SWOT component (Strengths, Weaknesses, Opportunities, or Threats).

    Attributes:
        step_name: Name of the SWOT analysis step being executed
        system_prompt: The system prompt to guide the language model
        sources: List of data sources to analyze
        output_model: The Pydantic model class for structured output
        company_name: The name of the company being analyzed
    """

    model_config = ConfigDict(frozen=True)

    component: SWOTComponent
    extraction_system_prompt: str
    sources: list[Source]
    extraction_output_model: type[SWOTExtractionModel]
    company_name: str


class ReportSummarizationContext(BaseModel):
    """
    Context information for summarizing SWOT analysis reports.

    Contains all the necessary information to summarize a SWOT analysis report.
    """

    model_config = ConfigDict(frozen=True)

    component: SWOTComponent
    summarization_system_prompt: str
    extraction_results: SWOTExtractionModel
    company_name: str


class ReportModificationContext(BaseModel):
    """
    Context information for modifying existing SWOT analysis reports.

    Contains the information needed to modify an already-generated SWOT analysis
    based on new sources or specific instructions.

    Attributes:
        step_name: Name of the SWOT analysis step being modified
        system_prompt: The system prompt to guide the language model
        modify_instruction: Specific instruction for how to modify the report
        structured_report: The existing report to be modified
        sources: List of new data sources to incorporate
    """

    model_config = ConfigDict(frozen=True)

    step_name: str
    system_prompt: str
    modify_instruction: str
    structured_report: SWOTExtractionModel
    sources: list[Source]
