import json
from typing import Any

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from dotenv import load_dotenv

from .state import AgentState, parameter_catalog

load_dotenv()


class ExtractedParameter(BaseModel):
    name: str = Field(
        description="Parameter name from the global catalog"
    )
    value: str | int | float | bool | None = Field(
        description="Extracted parameter value"
    )


class ExtractionResult(BaseModel):
    parameters: list[ExtractedParameter] = Field(
        default_factory=list
    )


llm = ChatOpenAI(model="gpt-4o")

extractor_llm = llm.with_structured_output(
    ExtractionResult
)


def extract_parameters(state: AgentState):

    system_prompt = f"""
    You are a parameter extraction agent.

    Extract parameters from the user conversation.

    Available parameter catalog:
    {json.dumps(parameter_catalog)}

    Rules:
    - Only extract parameters defined in the catalog.
    - Use conversation context to understand references.
    - Do not invent missing values.
    - Return only new or corrected values.
    - Do not decide which tool to execute.
    """

    result = extractor_llm.invoke([
        SystemMessage(content=system_prompt),
        *state["messages"]
    ])

    extracted = {
        param.name: param.value
        for param in result.parameters
        if param.name in parameter_catalog
    }

    return {
        "global_params": extracted
    }