from typing import Generator, Optional, Dict, Any, List, Union
from pydantic import Field, validator
from airtrain.core.skills import Skill, ProcessingError
from airtrain.core.schemas import InputSchema, OutputSchema
from .credentials import GroqCredentials
from .models_config import get_max_completion_tokens, get_model_config
from groq import Groq


class GroqInput(InputSchema):
    """Schema for Groq input"""

    user_input: str = Field(..., description="User's input text")
    system_prompt: str = Field(
        default="You are a helpful assistant.",
        description=(
            "System prompt to guide the model's behavior"
        ),
    )
    conversation_history: List[Dict[str, str]] = Field(
        default_factory=list,
        description=(
            "List of previous conversation messages in "
            "[{'role': 'user|assistant', 'content': 'message'}] format"
        ),
    )
    model: str = Field(
        default="llama-3.3-70b-versatile", 
        description="Groq model to use"
    )
    max_tokens: int = Field(
        default=4096, 
        description="Maximum tokens in response"
    )
    temperature: float = Field(
        default=0.7, 
        description="Temperature for response generation", 
        ge=0, 
        le=1
    )
    stream: bool = Field(
        default=False,
        description="Whether to stream the response progressively"
    )
    tools: Optional[List[Dict[str, Any]]] = Field(
        default=None,
        description=(
            "A list of tools the model may use. "
            "Currently only functions supported."
        ),
    )
    tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(
        default=None,
        description=(
            "Controls which tool is called by the model. "
            "'none', 'auto', or specific tool."
        ),
    )
    
    @validator('max_tokens')
    def validate_max_tokens(cls, v, values):
        """Validate that max_tokens doesn't exceed the model's limit."""
        if 'model' in values:
            model_id = values['model']
            max_limit = get_max_completion_tokens(model_id)
            if v > max_limit:
                return max_limit
        return v


class GroqOutput(OutputSchema):
    """Schema for Groq output"""

    response: str = Field(..., description="Model's response text")
    used_model: str = Field(..., description="Model used for generation")
    usage: Dict[str, Any] = Field(
        default_factory=dict, description="Usage statistics from the API"
    )
    tool_calls: Optional[List[Dict[str, Any]]] = Field(
        default=None, description="Tool calls generated by the model"
    )


class GroqChatSkill(Skill[GroqInput, GroqOutput]):
    """Skill for Groq chat"""

    input_schema = GroqInput
    output_schema = GroqOutput

    def __init__(self, credentials: Optional[GroqCredentials] = None):
        super().__init__()
        self.credentials = credentials or GroqCredentials.from_env()
        self.client = Groq(api_key=self.credentials.groq_api_key.get_secret_value())

    def _build_messages(self, input_data: GroqInput) -> List[Dict[str, str]]:
        """
        Build messages list from input data including conversation history.

        Args:
            input_data: The input data containing system prompt, conversation history, and user input

        Returns:
            List[Dict[str, str]]: List of messages in the format required by Groq
        """
        messages = [{"role": "system", "content": input_data.system_prompt}]

        # Add conversation history if present
        if input_data.conversation_history:
            messages.extend(input_data.conversation_history)

        # Add current user input
        messages.append({"role": "user", "content": input_data.user_input})

        return messages

    def process_stream(self, input_data: GroqInput) -> Generator[str, None, None]:
        """Process the input and stream the response token by token."""
        try:
            messages = self._build_messages(input_data)

            stream = self.client.chat.completions.create(
                model=input_data.model,
                messages=messages,
                temperature=input_data.temperature,
                max_tokens=input_data.max_tokens,
                stream=True,
            )

            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    yield chunk.choices[0].delta.content

        except Exception as e:
            raise ProcessingError(f"Groq streaming failed: {str(e)}")

    def process(self, input_data: GroqInput) -> GroqOutput:
        """Process the input and return the complete response."""
        try:
            if input_data.stream:
                response_chunks = []
                for chunk in self.process_stream(input_data):
                    response_chunks.append(chunk)
                response = "".join(response_chunks)
                usage = {}  # Usage stats not available in streaming
                tool_calls = None  # Tool calls not available in streaming
            else:
                messages = self._build_messages(input_data)
                
                # Prepare API call parameters
                api_params = {
                    "model": input_data.model,
                    "messages": messages,
                    "temperature": input_data.temperature,
                    "max_tokens": input_data.max_tokens,
                    "stream": False,
                }
                
                # Add tools and tool_choice if provided
                if input_data.tools:
                    api_params["tools"] = input_data.tools
                
                if input_data.tool_choice:
                    api_params["tool_choice"] = input_data.tool_choice
                
                completion = self.client.chat.completions.create(**api_params)
                response = completion.choices[0].message.content or ""
                
                # Extract usage information
                usage = {
                    "total_tokens": completion.usage.total_tokens,
                    "prompt_tokens": completion.usage.prompt_tokens,
                    "completion_tokens": completion.usage.completion_tokens,
                }
                
                # Check for tool calls in the response
                tool_calls = None
                if (
                    hasattr(completion.choices[0].message, "tool_calls") 
                    and completion.choices[0].message.tool_calls
                ):
                    tool_calls = [
                        {
                            "id": tool_call.id,
                            "type": tool_call.type,
                            "function": {
                                "name": tool_call.function.name,
                                "arguments": tool_call.function.arguments
                            }
                        }
                        for tool_call in completion.choices[0].message.tool_calls
                    ]

            return GroqOutput(
                response=response, 
                used_model=input_data.model, 
                usage=usage,
                tool_calls=tool_calls
            )

        except Exception as e:
            raise ProcessingError(f"Groq processing failed: {str(e)}")
