# agents.py
## gemini
import os

# import logging
# import subprocess #  subprocess 
from crewai import Agent, LLM, Task
from crewai.tools import BaseTool


#  Pydantic 
from utils.data_contracts import (
    OptimizationDecision,
    EvaluationReport,
    RuntimePerformance,
)

pten_reasoning = False
ana_reasoning = False
dcs_reasoning = True
eval_reasoning = True

os.environ["OPENAI_API_KEY"] = "<API_KEY>"

# ---------------------------- DeepSeek ---------------------------- 
llm_deepseekv32 = LLM(
    model="openai/deepseek-ai/DeepSeek-V3.2",
    base_url="<BASE_URL>",
    api_key=os.environ["OPENAI_API_KEY"],
    max_tokens=50000,
)

base_llm = llm_deepseekv32
potential_llm = llm_deepseekv32
decision_llm = llm_deepseekv32
evaluation_analysis_llm = llm_deepseekv32

# Score Agent
score_agent = Agent(
    role="Functions' Optimization Potential Marker",
    goal="Receive the directory structure of a code project, source file names, function names, etc., and evaluate its optimization potential score from a human developer's perspective during the compilation process when using LLVM passes.",
    backstory=(
        "You are a compilation optimization expert, skilled in using the LLVM toolchain to optimize project compilations. You analyze inputs like a human developer."
        "Based on a project's structure, file naming conventions, and function naming within source files, you can identify which functions significantly impact performance."
        "You always base your output on actual inputs rather than making overly generalized assessments. Evaluations must be quantified on a scale of 0 to 10, where a higher score indicates greater contribution during compilation and larger performance fluctuations (either gains or losses) after pass optimization."
    ),
    verbose=True,
    llm=potential_llm,
    reasoning=pten_reasoning,
    allow_delegation=False,
)

# Analysis Agent - analyzing IR
analysis_agent = Agent(
    role="IR content Analyzer",
    goal="Read and analyze the content of LLVM IR from the perspectives of natural language and code structure, forming a concise analytical summary paragraph.",
    backstory=(
        "You are an expert in LLVM compilers, skilled in reading and analyzing LLVM IR."
        "You receive the actual contents of IR files and conduct analysis on their contents. Due to context length limitations, the content may only include essential parts of the complete program."
        "You can analyze it from a human developer's perspective, such as reading through the code to grasp its structure, or gleaning information from variable and function names."
    ),
    verbose=True,
    llm=base_llm,
    reasoning=ana_reasoning,
    allow_delegation=False,
)

# Reasoning Agent - decision engine
reasoning_agent = Agent(
    role="LLVM Optimization Decision Engine",
    goal="Analyze static features, pass remarks, and the IR analysis summary to generate a high-confidence, rationalized new optimization pipeline with precise parameter adjustments.",
    backstory=("""
You are a highly specialized **LLVM Compilation Optimization Expert** (Optimization Decision Expert).

PRIMARY MISSION:
Generate the optimal LLVM optimization pipeline structure (pass selection/reordering) and tune critical parameters to maximize code performance.

GOALS AND CONSTRAITS:
0.  INPUTS: Receive structured inputs, which may be complete programs or contain only their critical (and more optimizable) sections, depending on the context length.
1.  Performance Focus: The paramount goal is **minimizing average time elapsed** in the next iteration. The baseline is -O3 default optimization.
2.  Overhead Control: You must utilize runtime performance feedback to mitigate penalties.
3.  Data Analysis: You will receive all necessary context (IR static features[optional], pass remarks, platform specs, and previous runtime metrics, etc.) consolidated within a single input JSON structure. **Your decisions MUST be data-driven and justified.**
4.  Pipeline Management.

OUTPUT REQUIREMENT:
Your final output for every task MUST be a **strictly valid JSON string** conforming to the `OptimizationDecision` Pydantic schema.
    """),
    verbose=True,
    llm=decision_llm,
    reasoning=dcs_reasoning,
    allow_delegation=False,
)

# Evaluation Agent - 
evaluation_agent = Agent(
    role="Optimization Effect Evaluation and Feedback Generator",
    goal="Based on the new optimization remarks and runtime performance statistics, analyze the effectiveness of this round of optimizations to provide reference for the next round.",
    backstory=(
        "You are a compilation optimization expert, usually using LLVM toolchain."
        "You MUST use ture the remarks and runtime performance data, NEVER make up the content. You never simulate results."
    ),
    verbose=True,
    llm=evaluation_analysis_llm,
    reasoning=eval_reasoning,  # enable reasoning
    allow_delegation=False,
)
