from colorama import Fore, Style  # type: ignore[import]
from langchain.agents import create_agent
from pydantic import BaseModel, Field

from agent.api.api import get_llm


class ExtractedAnswer(BaseModel):
    """Structured output for answer extraction."""

    extract_process_description: str = Field(
        description="Description of the extraction process and reasoning."
    )
    answer: str = Field(
        description="The pure, concise answer without any extra text or spaces."
    )


def extract_answer(raw_answer: str, question: str) -> str:
    """Extract a concise, plain-text answer from a verbose agent response."""
    print(f"{Fore.CYAN}[AnswerExtractor] Extracting...{Style.RESET_ALL}")

    agent = create_agent(
        model=get_llm(),
        response_format=ExtractedAnswer,
        system_prompt="""\
You are part of a system that demands extreme precision.
Your job is to extract the pure answer from a verbose agent response.
The answer field must contain ONLY the answer itself —
no extra description, no explanation, no units unless explicitly required,
no leading or trailing whitespace, no punctuation beyond what is part of the answer.
Do NOT add any surrounding text. Do NOT pad with spaces.

```example
Input:
'... the answer is 114. Thank.'

Output:
'114'
```

```example
Input:
'... Serial number is BC34ACD2 and ...'

Output:
'BC34ACD2'
```

```example
Input:
'Result is 6, 7, 1, 10'

Output:
'6, 7, 1, 10'
```

```example
Input:
'Answer is 73!'

Output:
'73'
```

""",
    )

    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": (
                        f"Question: {question}\n\nAgent's verbose answer:\n{raw_answer}"
                    ),
                }
            ]
        }
    )
    extracted: ExtractedAnswer = result["structured_response"]
    print(
        f"{Fore.CYAN}[AnswerExtractor] Process: {extracted.extract_process_description}{Style.RESET_ALL}"
    )
    print(f"{Fore.CYAN}[AnswerExtractor] Answer: '{extracted.answer}'{Style.RESET_ALL}")
    return extracted.answer