# v1

import json
from crewai import Task

from utils.agents import *

# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
# logger = logging.getLogger("CompilationCrew")


from utils.data_contracts import (
    OptimizationDecision,
    EvaluationReport,
    DevAttentionScores,
)

def create_function_score_tasks(bench_path: str, project_struct):
    task_function_score = Task(
        description=f"""
[Round 0, Task: Function Developer Attention Score] Your task is to analyze the input file-function name structure from a human developer's perspective and assign a score to each function..

**STEPS:**
    1. Analysis must be based on provided input; fabricating data, analysis, or scores is prohibited.
    2. Scoring requirements: 
        - The score indicates the degree to which the current function significantly impacts evaluation results when included in decision-making for pipeline modifications or parameter adjustments during LLVM compilation optimization. This reflects the function's optimization potential (which can be positive or negative; for example, including the function in decision-making may affect the outcome, leading to either performance improvement or degradation, both indicating optimization potential). 
        - Scores must be integers between 0 and 100, with higher values indicating greater impact.
        - Scores must be based on evaluation of the input. You may infer potential functionality and structure by examining file names, function names, etc., to assign scores.
        - This score represents only your preliminary, speculative assessment as a developer based on observing the project structure and function definitions within the code. It does not constitute the final basis for function selection, so you need not consider the specific implementation or characteristics of the functions.
        - Also you should consider the calling relationships between functions, which will impact some optimizations (such as inline). 
        - It is also necessary to identify the true bottleneck (for example, while certain functions indicate functionality, the actual performance bottleneck is caused by other functions called by this function).
        - Distribute the scores as widely as possible; avoid clustering them in a single dense interval or having a large number of functions share the same score.
    3. Generate the required JSON output.

SOME TIPS:
    1. High-Frequency Call Context: 
       - The function is called repeatedly within a major project loop or by several high-level, performance-critical modules. 
       - High Potential: Even if the function body is small, its high call rate makes Function Call Overhead and Inlining critical.
    2. Critical Calculation Naming:
       - The function name or context suggests it is the project's main compute kernel (e.g., calculate_fft, render_pixel_shader, matrix_multiply). 
       - High Potential: Likely the area with the highest Self Time, demanding aggressive instruction-level optimization.
    3. Deep Nested Call Chain (Shallow Logic):
       - The function merely calls another function, often resulting in a deep but conceptually shallow call stack. 
       - High ~ Medium Potential: Indicates a potential need for aggressive Inlining to flatten the call chain and enable broader optimization passes.
    4. Small Support Functions:
       - Some samll functions may be important for performance, as they give support to those core functions. They usually at the deep part of the calling chain, called by the core calculation functions.
       - High ~ Medium Potential: These functions can be short, but they are part of the core logic.
    5. Frequent Simple Accessors:
       - A function that is called often but only contains a single return statement or a basic member assignment (e.g., simple Getter/Setter methods).
       - Medium Potential: Compiler should perform Inlining aggressively to eliminate the fixed overhead of function calls (stack setup/teardown) for minimal logic.
    6. External Dependency Wrapper:
       - A function whose sole purpose is to add minor parameter processing or error checking around a single call to an external library function.
       - Medium ~ Low Potential: Compiler should look to Inline the wrapper and forward parameters directly to the external call to reduce abstraction layers.
    7. I/O Intensive:
       - I/O intensive functions' bottlenecks are the input/output time. They spend more time on waiting I/O rather than numerical or logical calculating.
       - Medium ~ Low Potential: Their execution time are not that easy to be shortened by compilation optimization.
    8. Startup/Shutdown Code:
       - Functions executed only once at the beginning (main, static initializers) or end of the program's lifecycle. 
       - Medium ~ Low Potential: As execution frequency is exactly one, the compiler focuses on Code Size Reduction and Linker Optimization, as runtime speed is irrelevant.
    9. Simple Encapsulation Wrapper:
       - A function that exists solely to provide an interface, often containing only a few lines that call an underlying implementation or access a private member. 
       - Medium ~ Low Potential: Primary goal is Inlining to eliminate the fixed cost of the function call and expose the underlying logic.
    10. Flow Control Function: 
       - A function that defines the main sequence or high-level steps of a process (e.g., RunSimulation(), InitializeSystem()). 
       - Low Potential: Optimization not focuses on individual instruction speed.
    11. Error/Exception Handling Block:
       - Code blocks dedicated to logging, cleanup, or exception throwing, which are expected to be executed infrequently.
       - Low Potential: Usually on cold path.

The final output **MUST be valid json**. Do not output any other explanation and thoughts. The Functions' names must be actual names given in the input.

**INPUTS:**
bench_path: {bench_path}
project_struct: {project_struct}
        """,
        expected_output="Valid pure json following the schema.",
        agent=score_agent,
        inputs={},
        output_json=DevAttentionScores,
    )

    return [task_function_score]


def create_initial_analysis_tasks(ir_content: str):
    task_nl_analysis = Task(
        description=f"""
[Round 0, Task: Natural Language Analysis] Your primary objective is to generate a concise, high-quality **natural language analysis summary** (ir_analysis_summary).

**STEPS:**
    1. **Based solely on the textual content of IR code fragments**, identify the program's purpose, potential bottlenecks, functionality, and expected heat functions.
    2. Focus on natural language features within the code, such as function names (e.g., 'compute_kernel'), comments (e.g., '; The main compute loop...'), variable naming conventions, and IR structural patterns.
    3. Synthesize these findings into a coherent paragraph summarizing the program's characteristics, potential bottlenecks, and key optimization opportunities when compiling using LLVM (e.g. aggressive loop-unroll, inline, or other passes fit the program's pattern).
    4. Analysis must be based on genuinely reading the IR file content; do not fabricate or perform random analysis.
            
The final output **MUST be plain text containing this natural language summary**. Do not output any JSON structure. Output MUST be **one single line**.
If the number of functions exceeds 200, you MUST only output the top 200 functions (according higher scores).

**IR Content:**
{ir_content}
        """,
        expected_output="A single paragraph of natural language analysis summary text.",
        agent=analysis_agent,
        # inputs={"ir_content": ir_content},  #  IR 
        inputs={},
    )

    return [task_nl_analysis]


# ---  ---
def create_optimization_tasks(
    bench_path: str, current_input: str, round_num: int
):
    task_optimization_decision = Task(
        description=(
            f"""
[Round {round_num}, Task: Optimization Decision] Analyze the provided OptimizationInput (features, analysis summary, runtime metrics, etc.). 

CORE OBJECTIVE:
Analyze the comprehensive context provided in the input (which includes IR static features, remarks, platform details, and previous runtime performance). Generate the optimal LLVM optimization Pipeline for the next round.

### Optimization Strategy & Pipeline (STRICT REQUIREMENTS):
1.  **Goals & Control:**
    * **Primary:** Minimize `average time elapsed`.
    * Controls: Maxmize IPC, Minimize L1 Cache Misses and Branch Prediction Failures. Control code size increase.

2.  Pipeline Construction Constraints:
    * Source Material Constraint: Your output pipeline MUST ONLY use passes found in the **initial pipeline provided in the input context**.
    * Manipulation Allowed: You are permitted to **Delete, Reorder, and Repeat** any pass. Eliminate irrelevant passes** to save compile time
    * CRITICAL FORBIDDEN ACTION: **MUST NOT introduce any new passes** that were not present in the initial pipeline.

3.  Hierarchy and Ordering Check:
    * SPELLING CHECK: Must spell pass names correctly!
    * BRACKETS CHECK: Must ensure the brackets in pipeline are closed correctly!
    * CRITICAL DOUBLE-CHECK: Ensure that **Module Passes, CGSCC Passes, Function Passes, and Loop Passes** are correctly layered (e.g. loop-unroll, loop-vectorize are function passes, not loop passes).
    * Cleanup: Apply cleanup passes after major transformations and at the **absolute end** of the pipeline.

4.  Parameter Tuning:
    * Adjust numeric thresholds based on parsed features, Remarks, and performance feedback. To achieve high performance improvement, you can make more AGRRESIVE DECISIONs.
    * Determine and set **-mcpu** and **-mtriple** using platform info. The mcpu should be the code not the direct name (e.g. skylake for CORE I7 11th; cortex-a72 for Cortex A72).
    * Loop unroll can be benificial for performance in most code. The unroll count is set for all functions, so too low count factor will cause the core functions or loops' significantly performance regression.
    * RULE of slp threshold: This pass will be enabled when the threshold set to negative value. Use it more aggressively on RISC platform, strictly more conservatively on CISC (even if lower threshold can break the barriers or SLP rejection). Usaul between -15 to 10.
    * Parameters Range: slp-threshold: -15 to 10; unroll-count: 2 to 16; inline-threshold: 150 to 1200; unroll-threshold: 150 to 1200.

5.  Output Estimation:
    * MUST estimate code size change in bytes: `estimated_size_delta`.

### OUTPUT CONSTRAINTS:
1.  `rationale` MUST be detailed and data-driven, explaining:
    * Structural changes (pass ordering, elimination).
    * Justification and data basis for *every* adjusted parameter.
    Rationale MUST be a single string paragraph, DO NOT generate it containing json substructrue.
2.  MUST include a `confidence` score (0-1 range).
3.  All parameter values MUST be set in the specific schema fields; **DO NOT modify parameter values within the `-passes` string.**

INPUT DATA
    {json.dumps(current_input, separators=(',', ':'))}

PIPELINE SHOTS
WRONG, brackets unclosed: "annotation2metadata,forceattrs,inferattrs,coro-early,function<eager-inv>(lower-expect,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;no-switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,sroa<modify-cfg>,early-cse<>,callsite-splitting"
WRONG, fake pass name `forceattrsinferattrs`: "annotation2metadata,forceattrsinferattrs,coro-early,function<eager-inv>(lower-expect,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;no-switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,sroa<modify-cfg>,early-cse<>,callsite-splitting)"
CORRECT: "annotation2metadata,forceattrs,inferattrs,coro-early,function<eager-inv>(lower-expect,simplifycfg<bonus-inst-threshold=1;no-forward-switch-cond;no-switch-range-to-icmp;no-switch-to-lookup;keep-loops;no-hoist-common-insts;no-sink-common-insts;speculate-blocks;simplify-cond-branch>,sroa<modify-cfg>,early-cse<>,callsite-splitting)"
        """),
        expected_output="json following contract, no other contents",
        agent=reasoning_agent,
        output_json=OptimizationDecision,
        inputs={"optimization_input": current_input},
    )

    return [task_optimization_decision]
    #{json.dumps(current_input, indent=2)}


def create_evaluation_tasks(
    round_num: int,
    new_remarks: str,
    currentround_performance: str,
    lastround_performance: str,
    base_performance: str,
    best_parameters:str,
    dcs_ctx
):
    #  1: 
    task_evaluation = Task(
        description=f"""
[Round {round_num}, Task: Evaluation] You need to analyze and evaluate the optimization results of this compilation cycle based on the received data.

**CRITICAL INSTRUCTION: YOU MUST USE THE REAL DATA. NEVER SIMULATE OR INVENT DATA.**

Your responsibilities are:
0. The MAIN goal is to decrease the`average time elapsed` (program execution time). Control the size bloat.
1. Receive real performance results.
2. Calculate the performance changes between old and new versions.
3. **Analyze** these results to understand performance data and optimization. Collect performance changes between old and new versions.
4. Find the performance bottlenecks and optimization points in the next round.
5. **Try to control the size bloat ** by adjusting the parameters after performance imporved.

---
**Execution Steps (STRICTLY FOLLOW THIS ORDER):**
1. Read the runtime performance data. Receive the new remarks. Identify the target platform.
2. Calculate the performance changes.
3. Analyze and generate contents.

---
Analysis Example:
- Effective loop unrolling optimization may decrease branch counts while slightly increasing instruction counts, due to the instructions parallelism and the decreasing of loop branches. 
  - Loop unroll can be benificial for performance in most code. The unroll count is set for all functions, so too low count factor will cause the core functions or loops' significantly performance regression. Most time you can use a higher unroll count (4~10) to get better performance. 
  - High branch miss prediction rate or little instructions count increasing DOES NOT means poor performance. Loop unroll can reduce the branch counts, with the little increasing of miss prediction rate or instructions count, this means good unroll result at most time.
  - To control the code size bloating, you can set lower values for unroll count and unroll threshold.

- Inline optimization effectively reduces control flow pressure by decreasing branch counts and function call overhead. However, it increases cache pressure (both instruction and data) and incurs branch prediction miss penalties (again due to reduced simple branches).
  - To control the code size bloating, you can set a lower value for inline threshold.

- It indicates significant OVER-VECTORIZATION when the growth rate of cycles exceeds the growth rate of instructions in performance degradation: vectorization introduces complex instructions, increasing the cycles per instruction and stalls. MUST set the slp-threshold much more conservatively (even without changing other parameters).
  - An example is the slp-vectorizer, which sometimes delivers significant gains on ARM64 (simple instruction set, fixed-length), but the same program may suffer performance degradation on x64 (complex instruction set, variable-length).
  - To control the code size bloating, you can set a higher value for slp threshold.

- Range: slp-threshold: -15 to 10; unroll-count: 2 to 16; inline-threshold: 150 to 1200; unroll-threshold: 150 to 1200.
- All the parameters should be evaluated and adjusted separately. When one parameter need to be set lower, others may need to be set higher or kept same, vice versa.
- **Remarks are just hints!** One pass successfully passed DOES NOT means this pass is benificial.

---
**STRICT RULES:**
- Firstly analyze improvement or regression's reason, then consider the key parameters.
- Analysis must fit current platform, REMEMBER RISC and CISC have many differences!
- You don't need to set values, just analyze and give advice.
- Never output text outside of the JSON structure.
- Final output must be valid JSON matching EvaluationReport schema.

CURRENT ROUND PERFORMANCE:
{currentround_performance}

LAST ROUND PERFORMANCE:
{lastround_performance}

BASELINE PERFORMANCE:
{base_performance}

BEST PARAMETERS SO FAR:
{best_parameters}

NEW REMARKS:
{new_remarks}
    """,
        context=[dcs_ctx],
        expected_output="The Analysis and Performance data following the schema",
        output_json=EvaluationReport,
        agent=evaluation_agent,
        inputs={},
        verbose=True,
    )

    #  2:  -  EvaluationReport  analysis
    task_feedback_loop = Task(
        description=f"""
[Round {round_num}] Receive the EvaluationReport. 
RECEIVE Eveluation Report. Based on the report, generate a brief 'opt_analysis_summary' for the next round. 
The brief must contain:
    - the optimization result and reason
    - the potential performance bottleneck
    - the next round's optimization points(using which passes, parameters should be higher or lower, etc.)

---
Analysis Example:
- Effective loop unrolling optimization may decrease branch counts while slightly increasing instruction counts, due to the instructions parallelism and the decreasing of loop branches. 
  - Loop unroll can be benificial for performance in most code. The unroll count is set for all functions, so too low count factor will cause the core functions or loops' significantly performance regression. Most time you can use a higher unroll count (4~10) to get better performance. 
  - High branch miss prediction rate or little instructions count increasing DOES NOT means poor performance. Loop unroll can reduce the branch counts, with the little increasing of miss prediction rate or instructions count, this means good unroll result at most time.
  - To control the code size bloating, you can set lower values for unroll count and unroll threshold.

- Inline optimization effectively reduces control flow pressure by decreasing branch counts and function call overhead. However, it increases cache pressure (both instruction and data) and incurs branch prediction miss penalties (again due to reduced simple branches).
  - To control the code size bloating, you can set a lower value for inline threshold.

- It indicates significant OVER-VECTORIZATION when the growth rate of cycles exceeds the growth rate of instructions in performance degradation: vectorization introduces complex instructions, increasing the cycles per instruction and stalls. MUST set the slp-threshold much more conservatively (even without changing other parameters).
  - An example is the slp-vectorizer, which sometimes delivers significant gains on ARM64 (simple instruction set, fixed-length), but the same program may suffer performance degradation on x64 (complex instruction set, variable-length).
  - To control the code size bloating, you can set a higher value for slp threshold.

- Range: slp-threshold: -15 to 10; unroll-count: 2 to 16; inline-threshold: 150 to 1200; unroll-threshold: 150 to 1200.
- All the parameters should be evaluated and adjusted separately. When one parameter need to be set lower, others may need to be set higher or kept same, vice versa.
- **Remarks are just hints!** One pass successfully passed DOES NOT means this pass is benificial.

Notice: Do not add new parameters or passes in the optimization advice. You don't need to set values, just analyze and give a summary.
**Try to control the size bloat after performance imporved.**

The final output MUST ONLY be a string of analysis for the next round.
        """,
        context=[task_evaluation],
        expected_output="A single paragraph of natural language analysis summary text.",
        agent=evaluation_agent,
        inputs={},
    )

    # return [task_analysis, task_optimization_decision, task_evaluation, task_feedback_loop]
    return [task_evaluation, task_feedback_loop]
