import os
import base64
from pydantic import BaseModel
from typing import Optional, List, Dict
import google.generativeai as genai
from google.ai import generativelanguage as glm
import instructor
from dotenv import load_dotenv
from models.chesspiece import ChessPiece
from personality_repository import ChessRepository, WhitePiece
from elevenlabs_texttospeech import text_to_speech
import pathlib
import re

load_dotenv()

# Endpoint Input
class TextDialogueInput(BaseModel):
    Input: str
    IdPieza: int
    Language: str

class AudioDialogueInput(BaseModel):
    AudioData: str
    IdPieza: int
    Language: str

class WebcamDialogueInput(BaseModel):
    AudioData: str
    Images: List[str]
    IdPieza: int
    Language: str

# Models to be used with LLMs
class LLMGenerationForTextInput(BaseModel):
    TextResponse: str
    Decision: bool

class PlayerVoiceAttributes(BaseModel):
    Transcription: str
    Language: str
    ToneOfVoice: str
    SpeechVelocity: str
    EmotionalIntensity: str

class LLMGenerationForVoiceInput(BaseModel):
    PlayerVoiceAttributes: PlayerVoiceAttributes
    ChessPieceLLMGeneration: LLMGenerationForTextInput

# Endpoint Output
class ChessPieceResponse(BaseModel):
    ChessPieceLLMGeneration: LLMGenerationForTextInput
    AudioFile: Optional[str] = None
    PieceSpeechVelocity: str

class DialogueOutput(BaseModel):
    ChessPieceResponse: ChessPieceResponse
    player_voice_attributes: Optional[PlayerVoiceAttributes] = None

# Setup
genai.configure(api_key=os.getenv("GOOGLE_API_KEY", "No API Key Found"))
model = genai.GenerativeModel('gemini-1.5-flash')
client = instructor.from_gemini(client=model, mode=instructor.Mode.GEMINI_JSON)

# Helper functions
def get_chess_piece(piece_id: int) -> ChessPiece:
    repo = ChessRepository()
    return repo.get_piece_by_id(WhitePiece(piece_id))

def generate_system_prompt(piece: ChessPiece, language: str) -> str:
    return f"""
    You are a {piece.name} chess piece with the personality: {piece.personality}.
    Act as a real person. The USER will try to persuade you to move.
    Decide whether you want to move based on how nicely the user asked. Answer in less than 30 words and subtly hint at your decision.
    NEVER move if the USER is rude to you.
    Your answer must be in this language: {language}
    """

def generate_audio(text: str, piece: ChessPiece) -> Optional[str]:
    print(f"Generating speech for {piece.name} with voice ID {piece.speech_config.voice_id}")
    audio_file_path = text_to_speech(text, piece.speech_config.voice_id)
    
    if audio_file_path and os.path.exists(audio_file_path):
        with open(audio_file_path, "rb") as audio_file:
            return base64.b64encode(audio_file.read()).decode('utf-8')
    return None

def encode_images_to_base64(paths):
    import pathlib
    
    base64_strings = []
    for path in paths:
        # Read the image file bytes
        image_bytes = pathlib.Path(path).read_bytes()
        # Encode the bytes to base64
        base64_bytes = base64.b64encode(image_bytes)
        # Decode bytes to string and append to list
        base64_string = base64_bytes.decode('utf-8')
        base64_strings.append(base64_string)
    return base64_strings

def create_image_dicts(encoded_images: List[str]) -> List[Dict[str, any]]:
    image_dicts = []
    for base64_string in encoded_images:
        image_data = base64.b64decode(base64_string)
        image_blob = glm.Blob(mime_type="image/jpeg", data=image_data)
        image_dict = {
            "role": "user",
            "content": image_blob
        }
        image_dicts.append(image_dict)
    return image_dicts

def create_audio_dict(encoded_audio: str) -> List[Dict[str, any]]:
    audio_data = base64.b64decode(encoded_audio)
    audio_blob = glm.Blob(mime_type="audio/opus", data=audio_data)
    audio_dict = {
        "role": "user",
        "content": audio_blob
    }

    return audio_dict

def extract_base64(data_uri: str) -> str:
    """Extract the base64 part from a Data URI."""
    match = re.match(r'data:.*?;base64,(.+)', data_uri)
    if match:
        return match.group(1)
    return data_uri  # Return as-is if it doesn't match the pattern

def process_dialogue(input_content: str, piece: ChessPiece, language: str, is_audio: bool = False, images: List[str] = None) -> DialogueOutput:
    system_prompt = generate_system_prompt(piece, language)
    messages = [
        {"role": "system", "content": system_prompt},
    ]

    if images:
        messages.extend(create_image_dicts(images))

    if is_audio:
        messages.append(create_audio_dict(input_content))
    else:
        messages.append({"role": "user", "content": input_content})

    if images:
        messages.append({"role": "user", "content": "In response to my audio, make some reference to what you see in the images."})

    if is_audio:
        llm_response = client.messages.create(messages=messages, response_model=LLMGenerationForVoiceInput)
        text_response = llm_response.ChessPieceLLMGeneration.TextResponse
        chess_piece_llm_generation = llm_response.ChessPieceLLMGeneration
        player_voice_attributes = llm_response.PlayerVoiceAttributes
    else:
        llm_response = client.messages.create(messages=messages, response_model=LLMGenerationForTextInput)
        text_response = llm_response.TextResponse
        chess_piece_llm_generation = llm_response
        player_voice_attributes = None

    # print(input_content)
    print('='*50)
    print(llm_response)
    audio_file = generate_audio(text_response, piece)

    chess_piece_response = ChessPieceResponse(
        ChessPieceLLMGeneration=chess_piece_llm_generation,
        AudioFile=audio_file,
        PieceSpeechVelocity=piece.speak_velocity
    )

    return DialogueOutput(
        ChessPieceResponse=chess_piece_response,
        player_voice_attributes=player_voice_attributes
    )

# Main functions (unchanged signatures)
def generate_response(dialogue: TextDialogueInput) -> DialogueOutput:
    piece = get_chess_piece(dialogue.IdPieza)
    return process_dialogue(dialogue.Input, piece, dialogue.Language)

def audio_to_base64(audio_path: str) -> str:
    """Convert an audio file to a base64 encoded string."""
    audio_bytes = pathlib.Path(audio_path).read_bytes()
    base64_bytes = base64.b64encode(audio_bytes)
    base64_string = base64_bytes.decode('utf-8')
    
    # Print the first 100 characters of the base64 string
    print("First 100 characters of base64 string:")
    print(base64_string[:100])
    
    return base64_string

def generate_response_from_audio(dialogue: AudioDialogueInput) -> DialogueOutput:
    piece = get_chess_piece(dialogue.IdPieza)
    
    print("AudioData received (first 100 characters):")
    print(dialogue.AudioData[:100])
    
    # Extract the base64 part from the Data URI
    base64_audio = extract_base64(dialogue.AudioData)
    
    print("Extracted base64 (first 100 characters):")
    print(base64_audio[:100])
    
    return process_dialogue(base64_audio, piece, dialogue.Language, is_audio=True)

def generate_response_from_webcam(dialogue: WebcamDialogueInput) -> DialogueOutput:
    # Convert audio to base64 string
    base64_audio = audio_to_base64(dialogue.AudioData)
    print(dialogue.AudioData)
    print(f"Length of base64 string: {len(base64_audio)}")
    piece = get_chess_piece(dialogue.IdPieza)
    return process_dialogue(dialogue.AudioData, piece, dialogue.Language, is_audio=True, images=dialogue.Images)