from enum import StrEnum

from openai.types.chat import ChatCompletionToolParam
from pydantic import BaseModel, ConfigDict, Field

from spark_advisor_models.model import AnalysisToolInput


class AgentToolName(StrEnum):
    GET_JOB_OVERVIEW = "get_job_overview"
    GET_STAGE_DETAILS = "get_stage_details"
    RUN_RULES_ENGINE = "run_rules_engine"
    CALCULATE_OPTIMAL_PARTITIONS = "calculate_optimal_partitions"
    COMPARE_CONFIGS = "compare_configs"
    SUBMIT_FINAL_REPORT = "submit_final_report"


class GetJobOverviewInput(BaseModel):
    """No input required — returns overview of the job being analyzed."""

    model_config = ConfigDict(frozen=True)


class GetStageDetailsInput(BaseModel):
    model_config = ConfigDict(frozen=True)

    stage_id: int = Field(description="Stage ID to retrieve detailed metrics for.")


class RunRulesEngineInput(BaseModel):
    """No input required — runs all 11 deterministic rules against the job."""
    model_config = ConfigDict(frozen=True)


class CalculateOptimalPartitionsInput(BaseModel):
    model_config = ConfigDict(frozen=True)

    total_shuffle_bytes: int = Field(description="Total shuffle bytes to partition.")
    target_partition_mb: int = Field(
        default=128,
        description="Target partition size in MB. Default 128MB.",
    )


class CompareConfigsInput(BaseModel):
    model_config = ConfigDict(frozen=True)

    proposed_changes: dict[str, str] = Field(
        description="Map of spark.* config keys to proposed new values.",
    )


def _tool(name: str, description: str, input_model: type[BaseModel]) -> ChatCompletionToolParam:
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description,
            "parameters": input_model.model_json_schema(),
        },
    }


AGENT_TOOLS: list[ChatCompletionToolParam] = [
    _tool(
        AgentToolName.GET_JOB_OVERVIEW,
        (
            "START HERE. Get a high-level summary of the Spark job: app ID, name, duration, "
            "stage count, executor count, total tasks, shuffle volumes, spill totals, and key "
            "spark.* configuration values (memory, cores, partitions, AQE, serializer). "
            "Also includes per-stage flags (SKEW, SPILL, GC, FAILURES) to quickly identify "
            "which stages need deeper investigation with get_stage_details. "
            "Call this first to understand the overall job health before diving into details."
        ),
        GetJobOverviewInput,
    ),
    _tool(
        AgentToolName.GET_STAGE_DETAILS,
        (
            "Get full metrics for a specific stage. Returns: task count, duration distribution "
            "(min/p25/median/p75/max), shuffle read/write bytes, spill to disk/memory, GC time, "
            "input/output bytes, failed task count, and complete task quantile distributions. "
            "Use this to investigate stages flagged with SKEW/SPILL/GC in the job overview. "
            "Compare min vs max task duration to quantify skew severity. "
            "Check if spill bytes are significant relative to shuffle read. "
            "Look at GC time as percentage of total task time (>20% is problematic)."
        ),
        GetStageDetailsInput,
    ),
    _tool(
        AgentToolName.RUN_RULES_ENGINE,
        (
            "Run the deterministic rules engine (17 rules covering: data skew, spill to disk, "
            "GC pressure, shuffle partition sizing, executor idle/underutilization, task failures, "
            "small files, broadcast join threshold, serializer choice, dynamic allocation, "
            "memory overhead, driver memory, AQE enablement, excessive stages, shuffle data volume, "
            "and input data skew). Returns a list of findings with severity (CRITICAL/WARNING/INFO), "
            "current values, recommended values, and estimated impact. "
            "Use rule findings to validate your own observations and discover issues you may have missed. "
            "Rule recommendations provide concrete spark.* config values you can include in your report."
        ),
        RunRulesEngineInput,
    ),
    _tool(
        AgentToolName.CALCULATE_OPTIMAL_PARTITIONS,
        (
            "Calculate the optimal spark.sql.shuffle.partitions value given total shuffle bytes "
            "and a target partition size (default 128MB — the Spark golden rule). "
            "Returns: current partition count, optimal count, and actual partition size. "
            "Use this when you see shuffle stages with partitions far from 128MB target, "
            "or when the rules engine flags ShufflePartitionsRule. "
            "Pass the total shuffle bytes from the stage with the largest shuffle volume."
        ),
        CalculateOptimalPartitionsInput,
    ),
    _tool(
        AgentToolName.COMPARE_CONFIGS,
        (
            "Compare your proposed Spark config changes against the job's current values. "
            "Shows a side-by-side diff of current vs proposed for each parameter. "
            "Use this BEFORE submitting the final report to verify your recommendations are correct — "
            "check that you're not recommending values the job already uses, "
            "that numeric changes go in the right direction, and that all recommendations are concrete "
            "(specific values, not vague guidance). This is your pre-flight check."
        ),
        CompareConfigsInput,
    ),
    _tool(
        AgentToolName.SUBMIT_FINAL_REPORT,
        (
            "Submit the final analysis report. This terminates the agent loop. "
            "Must include: a 1-2 sentence summary of the job's health, overall severity "
            "(critical/warning/info), up to 7 prioritized recommendations with concrete "
            "spark.* config values (not vague advice), estimated impact for each, and a causal chain "
            "explaining how problems are connected (e.g., skew → spill → GC → straggler). "
            "Only call this after you have gathered enough data to make confident recommendations."
        ),
        AnalysisToolInput,
    ),
]
