# v1

from pydantic import BaseModel, Field, RootModel
from typing import List, Dict, Union, Optional

# 0.


class IRContent(BaseModel):
    content: str = Field(description="IR ")


class FunctionCall(BaseModel):
    name: str
    callee: List[str] = Field(
        description="Functions directly called by current function.")


class SourceStructure(BaseModel):
    """Structure of Single Source File"""
    filename: str = Field(description="Source File Path and Name")
    functions: List[FunctionCall] = Field(
        description="Functions of Current Source")


class ProjectStructure(BaseModel):
    """Structure of Project Dir"""
    bench_path: str
    proj_struct: List[SourceStructure]


class FunctionAttentionScore(BaseModel):
    name: str = Field(description="function name")
    score: float = Field(
        description="Developers recognize the importance of this function in compilation optimization, specifically its potential for impacting the optimization.")


class DevAttentionScores(BaseModel):
    bench_path: str
    scores: List[FunctionAttentionScore]
    rationale: str = Field(
        description="Detailed explanation of the scores."
    )


# 1. LLVM Pass Remarks
class PassRemarkDetail(BaseModel):
    """Structure of a Single Pass Remarks"""

    remark_order: List[int] = Field(
        description="List of remarks recorded in the order they were generated"
    )
    count: int = Field(description="Number of occurrences")
    Name: str = Field(description="Pass name")
    ReasonString: str = Field(description="Reason Description")


class PassRemarkDetailNoOrder(BaseModel):
    """Structure of a Single Pass Remarks (without remark position order)"""

    count: int = Field(description="Number of occurrences")
    Name: str = Field(description="Pass name")
    ReasonString: str = Field(description="Reason Description")


class FunctionRemarks(BaseModel):
    """Summary of Optimization Notes for Individual Functions"""

    analysis: Dict[str, List[PassRemarkDetail]]
    missed: Dict[str, List[PassRemarkDetail]]
    passed: Dict[str, List[PassRemarkDetail]]


class FunctionRemarksNoOrder(BaseModel):
    """Summary of Optimization Notes for Individual Functions (without remark position order)"""

    analysis: Dict[str, List[PassRemarkDetailNoOrder]]
    missed: Dict[str, List[PassRemarkDetailNoOrder]]
    passed: Dict[str, List[PassRemarkDetailNoOrder]]


# 2. IR
class FunctionFeatures(BaseModel):
    """Function-level detailed static attributes"""

    Blocks: int
    Loops: int
    Calls: int
    ConditionalBranch: int


class GenericFeatures(BaseModel):
    """General or auxiliary static features"""

    is_indirect: bool
    num_loops: int
    cost_estimate: float


class FeaturesUnion(BaseModel):
    callee: FunctionFeatures = None
    caller: FunctionFeatures = None
    generic: GenericFeatures = None
    # generic: Optional[GenericFeatures] = None


class StaticFeatAndRemarksItem(BaseModel):
    """Static characteristics and remarks for individual functions"""

    name: str
    features: FeaturesUnion = None
    # remarks: FunctionRemarks = None
    remarks: FunctionRemarksNoOrder = None


class Remarks(BaseModel):
    """Remarks for individual functions"""

    name: str
    remarks: Optional[FunctionRemarksNoOrder] = None


# 3.
class PlatformInfo(BaseModel):
    """Compilation Environment and Target CPU Architecture Information"""

    target_arch: str = Field(description="e.g., x86_64, aarch64")
    compiler: str = Field(description="e.g., clang 17.0.6")
    cpu: str = Field(description="e.g., Intel I7-10700K")
    freq: str = Field(description="CPU Frequency")
    L1dcache: str
    L1icache: str
    L2cache: str
    L3cache: str


class RuntimePerformance(BaseModel):
    """Runtime performance data collected using `perf stat`"""

    file_size_bytes: int = Field(description="Compiled File Size (Bytes)")
    average_cycles: float
    average_instructions: float
    average_uops_issued_any: Optional[float] = Field(
        default=None, alias="average_uops_issued_any")
    average_mem_load_retired_l1_miss: Optional[float] = Field(
        default=None, alias="average_mem_load_retired_l1_miss"
    )
    average_mem_load_retired_l2_miss: Optional[float] = Field(
        default=None, alias="average_mem_load_retired_l2_miss"
    )
    average_mem_load_retired_l3_miss: Optional[float] = Field(
        default=None, alias="average_mem_load_retired_l3_miss"
    )
    average_cycle_activity_stalls_mem_any: Optional[float] = Field(
        default=None, alias="average_cycle_activity_stalls_mem_any"
    )
    average_L1_icache_load_misses: Optional[float] = Field(
        default=None, alias="average_L1_icache_load_misses"
    )
    average_stalled_cycles_backend: Optional[float] = Field(
        default=None, alias="average_stalled_cycles_backend"
    )
    average_stalled_cycles_frontend: Optional[float] = Field(
        default=None, alias="average_stalled_cycles_frontend"
    )
    average_branches: float
    average_branch_misses: float = Field(..., alias="average_branch_misses")
    average_time_elapsed: float
    average_time_user: float
    average_time_sys: float
    IPC: float = Field(description="Instructions Per Cycle")
    BR_MISP: float = Field(description="Branch Miss Predict Rate")


class OptimizationInput(BaseModel):
    """Core Input Data Contract for the Reasoning Agent"""

    ir_analysis_summary: str = Field(
        description="Analysis Agent generates natural language observations and optimization opportunity summaries based on IR static features and annotations."
    )
    static_feat_and_remarks: List[StaticFeatAndRemarksItem]
    platform_info: PlatformInfo
    runtime_performance: Optional[RuntimePerformance] = Field(
        None, description="Runtime Performance Statistics"
    )
    pipeline_params_to_optimize: Optional[Dict[str, Union[str, int, float]]] = Field(
        None, description="Pipelines and parameters to be optimized and adjusted"
    )


class OptimizationInputFeedback(BaseModel):
    """Core Input Data Contract for the Reasoning Agent (Feedback)"""

    opt_analysis_summary: str = Field(
        description="Evaluation Agent generates a summary of analysis and optimization opportunities based on optimization effectiveness and LLVM-generated optimization remarks."
    )
    new_remarks: List[Remarks]
    runtime_performance: Optional[RuntimePerformance] = Field(
        None, description="Runtime Performance Statistics"
    )
    previous_pipeline_params: Optional[Dict[str, Union[str, int, float]]] = Field(
        # previous_pipeline_params: Optional[OptimizationDecision] = Field(
        None,
        description="The pipeline and parameters adopted in this round of optimization",
    )


# ====================================================================
# Reasoning Agent  / Evaluation Agent
# ====================================================================


class PipelineParameterAdjustments(BaseModel):
    """Key Pass Parameter Values for LLM Prediction"""

    optimized_pipeline: str = Field(
        description="New or adjusted LLVM pass pipeline strings, limited to rearranging pass order or removing passes. **Do not add new passes or new embedded parameters**."
    )
    unroll_count: Union[int, str] = Field(
        ..., alias="unroll-count", description="Loop unroll factor, int or 'auto'"
    )
    unroll_allow_partial: Union[bool, str] = Field(
        ...,
        alias="unroll-allow-partial",
        description="Whether partial unroll is allowed, bool or 'auto'",
    )
    unroll_threshold: Union[int, str] = Field(
        ..., alias="unroll-threshold", description="Loop unroll threshold, int or 'auto'"
    )
    inline_threshold: Union[int, str] = Field(
        ..., alias="inline-threshold", description="inline threshold, int or 'auto'"
    )
    slp_threshold: Union[int, str] = Field(
        ...,
        alias="slp-threshold",
        description="SLP vectorization threshold; negative values indicate enabled, int  or 'auto'",
    )
    devirt_param: Union[int, str] = Field(
        ..., alias="devirt<>", description="Dereferencing virtual function parameters"
    )
    mtriple: str = Field(
        description="Target Machine Triple (e.g., x86_64-unknown-linux-gnu)"
    )
    mcpu: str = Field(description="Target CPU Type (e.g., skylake, native)")


class OptimizationDecision(BaseModel):
    """Reasoning Agent Output Data Contract"""

    passes_param_adjustments: PipelineParameterAdjustments
    confidence: float = Field(
        description="Confidence Index for This Optimization Decision (0.0 - 1.0)"
    )
    rationale: str = Field(
        description="Detailed optimization choices and the reasoning process behind parameter predictions"
    )
    similarity_vs_init_pipeline: float = Field(
        ...,
        alias="similarity vs init pipeline",
        description="Similarity to the pipeline before adjustment (0.0 - 1.0)",
    )


# ====================================================================
# Evaluation Agent  / Reasoning Agent (Feedback)
# ====================================================================


class EvaluationAnalysis(BaseModel):
    """Text Analysis Section of the Performance Evaluation Report"""

    analysis: str = Field(
        description="A brief textual analysis of the optimization results (e.g., optimization potential, bottlenecks)"
    )
    performance: str = Field(
        description="Summary of performance results (e.g., '1.xx speedup ratio compared to baseline')"
    )
    llvm_diagnostics: str = Field(
        ...,
        alias="LLVM diagnostics",
        description="Compiler or LLVM runtime diagnostic information",
    )


class EvaluationReport(BaseModel):
    """Output Data Contract for the Evaluation Agent"""

    analysis_report: EvaluationAnalysis
    # runtime_performance: RuntimePerformance
    # new_remarks: List[Remarks] = Field(
    #    description="Following the optimization of this LLM prompt, the newly generated Pass Remarks (for use in the next round of decision-making)"
    # )


##### ----- Pydantic  ----- #####
# json :
# json.dumps(json_dict)
# json.loads(json_str)

#  json  model :
# model = EvaluationReport.model_validate_json(str)

#  model :
# next_input_model = OptimizationInputFeedback(
#     opt_analysis_summary=optimization_analysis.raw,
#     new_remarks=remarks,
#     runtime_performance=json.loads(performance),  #  type  Dict,  dict
#     previous_pipeline_params=json.loads(optimization_decision)["passes_param_adjustments"],
# )

#  json str:
# model_json = model.json()

# :
# task_result = crew.kickoff()
# result_json = task_result.json_dict
# result_str = task_result.raw  # .raw.strip()
