from typing import List, Optional, Dict, Any, Generator, Union
from pydantic import Field
from openai import OpenAI

from airtrain.core.skills import Skill, ProcessingError
from airtrain.core.schemas import InputSchema, OutputSchema
from .credentials import FireworksCredentials


class FireworksInput(InputSchema):
    """Schema for Fireworks AI chat 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, Any]] = Field(
        default_factory=list,
        description="List of previous conversation messages",
    )
    model: str = Field(
        default="accounts/fireworks/models/deepseek-r1",
        description="Fireworks AI model to use",
    )
    temperature: float = Field(
        default=0.7, description="Temperature for response generation", ge=0, le=1
    )
    max_tokens: Optional[int] = Field(
        default=131072, description="Maximum tokens in response"
    )
    context_length_exceeded_behavior: str = Field(
        default="truncate", description="Behavior when context length is exceeded"
    )
    stream: bool = Field(
        default=False,
        description="Whether to stream the response token by token",
    )
    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."
        ),
    )


class FireworksOutput(OutputSchema):
    """Schema for Fireworks AI chat output"""

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


class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
    """Skill for interacting with Fireworks AI models"""

    input_schema = FireworksInput
    output_schema = FireworksOutput

    def __init__(self, credentials: Optional[FireworksCredentials] = None):
        """Initialize the skill with optional credentials"""
        super().__init__()
        self.credentials = credentials or FireworksCredentials.from_env()
        self.client = OpenAI(
            base_url="https://api.fireworks.ai/inference/v1",
            api_key=self.credentials.fireworks_api_key.get_secret_value(),
        )

    def _build_messages(self, input_data: FireworksInput) -> List[Dict[str, Any]]:
        """Build messages list from input data including conversation history."""
        messages = [{"role": "system", "content": input_data.system_prompt}]

        if input_data.conversation_history:
            messages.extend(input_data.conversation_history)

        messages.append({"role": "user", "content": input_data.user_input})
        return messages

    def process_stream(self, input_data: FireworksInput) -> 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"Fireworks streaming failed: {str(e)}")

    def process(self, input_data: FireworksInput) -> FireworksOutput:
        """Process the input and return the complete response."""
        try:
            if input_data.stream:
                # For streaming, collect the entire response
                response_chunks = []
                for chunk in self.process_stream(input_data):
                    response_chunks.append(chunk)
                response = "".join(response_chunks)
                
                # Create completion object for usage stats
                messages = self._build_messages(input_data)
                completion = self.client.chat.completions.create(
                    model=input_data.model,
                    messages=messages,
                    temperature=input_data.temperature,
                    max_tokens=input_data.max_tokens,
                    stream=False,
                )
            else:
                # For non-streaming, use regular completion
                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 ""
            
            # 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 FireworksOutput(
                response=response,
                used_model=input_data.model,
                usage={
                    "total_tokens": completion.usage.total_tokens,
                    "prompt_tokens": completion.usage.prompt_tokens,
                    "completion_tokens": completion.usage.completion_tokens,
                },
                tool_calls=tool_calls
            )

        except Exception as e:
            raise ProcessingError(f"Fireworks chat failed: {str(e)}")
