from typing import Dict, List, Optional, TypedDict, Literal, TypeVar, Generic
from pydantic import BaseModel, Field, constr
from typing import Union
T = TypeVar('T')

class BaseOutputList(BaseModel, Generic[T]):
    items: List[T]
    
    def to_dict(self) -> dict:
        return self.model_dump()
    
class PatientCase(TypedDict):
    patient_info: Optional[str]
    partial_patient_info: Optional[str]
    options: List[Dict[str, float]]
    ground_truth: Optional[str]
    metadata: Optional[Dict[str, Union[str, int, float]]]  # New field for metadata

    @classmethod
    def create(cls, patient_info: Optional[str] = None, partial_patient_info: Optional[str] = None, options=None, ground_truth: Optional[str] = None, metadata: Optional[Dict[str, Union[str, int, float]]] = None) -> 'PatientCase':
        """
        Create a new PatientCase with default values.

        Args:
            patient_info: The full patient information
            partial_patient_info: The partial patient information
            ground_truth: The ground truth diagnosis

        Returns:
            A new PatientCase instance with initialized fields
        """
        return {
            "patient_info": patient_info if patient_info is not None else "",
            "partial_patient_info": partial_patient_info if partial_patient_info is not None else "",
            "options": options,
            "ground_truth": ground_truth if ground_truth is not None else "",
            "metadata": metadata if metadata is not None else {}
        }


class QA(TypedDict):
    """
    History represents the history of the patient.
    """
    question: str
    answer: str

    @classmethod
    def create(cls, question: Optional[str] = None, answer: Optional[str] = None) -> 'QA':
        """
        Create a new QA with default values.

        Args:
            question: The question to ask
            answer: The answer to the question

        Returns:
            A new QA instance with initialized fields
        """
        return {
            "question": question,
            "answer": answer
        }
"""
PossibleDisease represents a possible disease that the patient may have. 0.9-1.0, if the disease is expected to be the correct diagnosis, 0.6-0.9 if it is a possible diagnosis, and below 0.6 if it is unlikely.
"""

class PossibleDisease(BaseModel):
    disease: str
    confidence: float

'''class PossibleDisease(BaseModel):
    """
    PossibleDisease represents a possible disease that the patient may have.

    Attributes:
        disease (str): The name of the disease.
        probability (float): The probability of the disease being the correct diagnosis.
        rationale (str): The reasoning behind the diagnosis.
    """
    disease: str
    probability: float
    rationale: str'''
    
'''chain_of_thoughts: str = Field(
max_length=1500,
description="Reasoning steps for the diagnosis.",
#example="3yo with acute jaundice, dark urine, and recent URI. Labs show hemolysis (↑LDH, ↓haptoglobin) and C3d+ DAT. Key findings: 1) No hepatosplenomegaly → rules out chronic liver disease. 2) C3d dominance + post-URI onset → suggests cold antibody process. Most likely: Paroxysmal Cold Hemoglobinuria (PCH), with CAD as alternate."
)
diagnosis: List[PossibleDisease] = Field(
    min_length=5,
    max_length=8,
    description="List of possible diseases with confidence scores"
)'''

# class DifferentialDiagnosis(BaseModel):
#     diagnosis: List[PossibleDisease] = Field(
#         min_length=4,
#         max_length=8,
#         description="List of possible diseases with confidence scores"
#     )
    
#     def to_dict(self) -> dict:
#         """Convert to dictionary"""
#         return self.model_dump()
    
class MentalCTScan(BaseModel):
    """
    MentalCTScan represents a single differential diagnosis based on anatomical layer.
    """
    layer: Literal["Anterior", "Mid-layer", "Posterior", "Vascular"]
    disease: str
    confidence: float = Field(ge=0.0, le=1.0)

class MentalCTScanDDx(BaseModel):
    """
    Differential diagnoses based on anatomical differential approach.
    """
    diagnosis: List[MentalCTScan] = Field(
        min_length=10,
        max_length=20,
        description="List of anatomical layer-based differential diagnoses"
    )

    def to_dict(self) -> dict:
        """Convert to dictionary"""
        return self.model_dump()

class SingleCategory(BaseModel):
        """
        Classify the question into one category.
        """
        question: str
        category: Literal["Lab/Imaging", "Symptom", "Physical Exam", "Social History", "Past Medical History", "Review of Systems", "Demographics"]

class Category(BaseModel):
    categories: List[SingleCategory]

    def to_dict(self) -> dict:
        """Convert to dictionary"""
        return self.model_dump().get("categories", [])

class DiseaseLocation(BaseModel):
    """
    DiseaseLocation represents the location of the disease in the body.

    Attributes:
        location (str): The anatomical location of the disease.
        potential_disease (List[str]): List of 5 potential diseases associated with the location. The diseases should be broader categories, not specific diseases, i.e.: Pulmonary nodule (broad category), Primary lung adenocarcinoma (specific disease).
        supporting_evidence (str): Evidence supporting the location.
        rule_out_test (str): Test to rule out other locations.
    """
    location: str
    potential_disease: List[str]
    supporting_evidence: str
    rule_out_test: str

class DiseaseLocationOutput(BaseModel):
    items: List[DiseaseLocation]
    
    def to_dict(self) -> dict:
        return self.model_dump()

class DifferentialDiagnosis(BaseModel):
    items: List[PossibleDisease]
    
    def to_dict(self) -> dict:
        return self.model_dump()

class AgentState(TypedDict):
    """
    AgentState represents the current state of the agent.
    """
    num_of_interactions: int
    patient_case: PatientCase
    differential_diagnosis: Optional[Union[DifferentialDiagnosis, MentalCTScanDDx]]
    inquiries_history: List[Optional[QA]]
    questions_history: List[str]

    @classmethod
    def create(cls, patient_case: PatientCase) -> 'AgentState':
        """
        Create a new AgentState with default values.
        
        Args:
            patient_case: The patient case information
            
        Returns:
            A new        return self.model_dump().get("categories", [])ized fields
        """
        return {
            "num_of_interactions": 0,
            "patient_case": patient_case,
            "differential_diagnosis": None,
            "inquiries_history": [],
            "questions_history": []
        }
