from typing import Type, TypeVar, Optional, List, Dict, Any, Generator
from pydantic import BaseModel, Field
import requests
import json
from loguru import logger
import re

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

ResponseT = TypeVar("ResponseT", bound=BaseModel)


class FireworksStructuredRequestInput(InputSchema):
    """Schema for Fireworks AI structured output input using requests"""

    user_input: str = Field(..., description="User's input text")
    system_prompt: str = Field(
        default="You are a helpful assistant that provides structured data.",
        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",
    )
    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: int = Field(default=4096, description="Maximum tokens in response")
    response_model: Type[ResponseT]
    stream: bool = Field(
        default=False,
        description="Whether to stream the response",
    )

    class Config:
        arbitrary_types_allowed = True


class FireworksStructuredRequestOutput(OutputSchema):
    """Schema for Fireworks AI structured output"""

    parsed_response: Any
    used_model: str
    usage: Dict[str, int]
    reasoning: Optional[str] = None


class FireworksStructuredRequestSkill(
    Skill[FireworksStructuredRequestInput, FireworksStructuredRequestOutput]
):
    """Skill for getting structured responses from Fireworks AI using requests"""

    input_schema = FireworksStructuredRequestInput
    output_schema = FireworksStructuredRequestOutput
    BASE_URL = "https://api.fireworks.ai/inference/v1/chat/completions"

    def __init__(self, credentials: Optional[FireworksCredentials] = None):
        """Initialize the skill with optional credentials"""
        super().__init__()
        self.credentials = credentials or FireworksCredentials.from_env()
        self.headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
        }

    def _build_messages(
        self, input_data: FireworksStructuredRequestInput
    ) -> List[Dict[str, str]]:
        """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 _build_payload(
        self, input_data: FireworksStructuredRequestInput
    ) -> Dict[str, Any]:
        """Build the request payload."""
        return {
            "model": input_data.model,
            "messages": self._build_messages(input_data),
            "temperature": input_data.temperature,
            "max_tokens": input_data.max_tokens,
            "stream": input_data.stream,
            "response_format": {
                "type": "json_object",
                "schema": {
                    **input_data.response_model.model_json_schema(),
                    "required": [
                        field
                        for field, _ in input_data.response_model.model_fields.items()
                    ],
                },
            },
        }

    def process_stream(
        self, input_data: FireworksStructuredRequestInput
    ) -> Generator[Dict[str, Any], None, None]:
        """Process the input and stream the response."""
        try:
            payload = self._build_payload(input_data)
            response = requests.post(
                self.BASE_URL,
                headers=self.headers,
                data=json.dumps(payload),
                stream=True,
            )
            response.raise_for_status()

            json_buffer = []
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode("utf-8").removeprefix("data: "))
                        if data["choices"][0]["delta"].get("content"):
                            content = data["choices"][0]["delta"]["content"]
                            json_buffer.append(content)
                            yield {"chunk": content}
                    except json.JSONDecodeError:
                        continue

            # Once complete, parse the full response with think tags
            complete_response = "".join(json_buffer)
            reasoning, json_str = self._parse_response_content(complete_response)

            try:
                parsed_response = input_data.response_model.model_validate_json(
                    json_str
                )
                yield {"complete": parsed_response, "reasoning": reasoning}
            except Exception as e:
                raise ProcessingError(f"Failed to parse JSON response: {str(e)}")

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

    def _parse_response_content(self, content: str) -> tuple[Optional[str], str]:
        """Parse response content to extract reasoning and JSON."""
        # Extract reasoning if present
        reasoning_match = re.search(r"<think>(.*?)</think>", content, re.DOTALL)
        reasoning = reasoning_match.group(1).strip() if reasoning_match else None

        # Extract JSON
        json_match = re.search(r"</think>\s*(\{.*\})", content, re.DOTALL)
        json_str = json_match.group(1).strip() if json_match else content

        return reasoning, json_str

    def process(
        self, input_data: FireworksStructuredRequestInput
    ) -> FireworksStructuredRequestOutput:
        """Process the input and return structured response."""
        try:
            if input_data.stream:
                # For streaming, collect and parse the entire response
                json_buffer = []
                parsed_response = None
                reasoning = None

                for chunk in self.process_stream(input_data):
                    if "chunk" in chunk:
                        json_buffer.append(chunk["chunk"])
                    elif "complete" in chunk:
                        parsed_response = chunk["complete"]
                        reasoning = chunk.get("reasoning")

                if parsed_response is None:
                    raise ProcessingError("Failed to parse streamed response")

                return FireworksStructuredRequestOutput(
                    parsed_response=parsed_response,
                    used_model=input_data.model,
                    usage={},  # Usage stats not available in streaming mode
                    reasoning=reasoning,
                )
            else:
                # For non-streaming, use regular request
                payload = self._build_payload(input_data)
                response = requests.post(
                    self.BASE_URL, headers=self.headers, data=json.dumps(payload)
                )
                response.raise_for_status()
                data = response.json()

                response_content = data["choices"][0]["message"]["content"]

                # Parse the response content to extract reasoning and JSON
                reasoning, json_str = self._parse_response_content(response_content)

                # Parse the JSON string into the specified model
                parsed_response = input_data.response_model.model_validate_json(
                    json_str
                )

                return FireworksStructuredRequestOutput(
                    parsed_response=parsed_response,
                    used_model=input_data.model,
                    usage=data["usage"],
                    reasoning=reasoning,  # Add reasoning to output if present
                )

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