import io
import os
import requests
import subprocess
import tempfile
import pandas as pd
import whisper
import warnings
import numpy as np
from typing import Literal
from pydub import AudioSegment
from smolagents import tool
import urllib.parse
from loguru import logger

warnings.filterwarnings("ignore")

whisperModel = Literal['base', 'small', 'tiny']


@tool
def get_files_by_task_id(task_id: str) -> bytes:
    """
    Download a file associated with a Hugging Face course question by task ID.

    Args:
        task_id (str): The unique identifier for the question/task whose file should be downloaded.

    Returns:
        bytes: The binary content of the downloaded file (e.g., Excel, image, etc.).

    Raises:
        requests.HTTPError: If the HTTP request fails (non-200 status code).
        requests.RequestException: For network-related errors.

    Example:
        >>> file_bytes = get_files('7bd855d8-463d-4ed5-93ca-5fe35145f733')
        >>> with open('output.xlsx', 'wb') as f:
        ...     f.write(file_bytes)
    """
    url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
    headers = {"accept": "application/json"}
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        return response.content
    except requests.HTTPError as http_err:
        raise requests.HTTPError(f"HTTP error occurred while fetching file for task_id '{task_id}': {http_err}")
    except requests.RequestException as req_err:
        raise requests.RequestException(
            f"Network error occurred while fetching file for task_id '{task_id}': {req_err}")


@tool
def excel_bytes_to_dataframe(excel_bytes: bytes) -> pd.DataFrame:
    """
    Convert Excel file bytes to a pandas DataFrame.

    Args:
        excel_bytes (bytes): The binary content of an Excel file.

    Returns:
        pd.DataFrame: A pandas DataFrame containing the Excel data.

    Raises:
        ValueError: If the bytes cannot be parsed as a valid Excel file.
    """
    try:
        return pd.read_excel(io.BytesIO(excel_bytes))
    except Exception as e:
        raise ValueError(f"Failed to parse Excel bytes: {str(e)}")


@tool
def csv_bytes_to_dataframe(csv_bytes: bytes, encoding: str = 'utf-8') -> pd.DataFrame:
    """
    Convert CSV file bytes to a pandas DataFrame.

    Args:
        csv_bytes (bytes): The binary content of a CSV file.
        encoding (str): The encoding of the CSV file (default 'utf-8').

    Returns:
        pd.DataFrame: A pandas DataFrame containing the CSV data.

    Raises:
        ValueError: If the bytes cannot be parsed as a valid CSV file.
    """
    try:
        return pd.read_csv(io.BytesIO(csv_bytes), encoding=encoding)
    except Exception as e:
        raise ValueError(f"Failed to parse CSV bytes: {str(e)}")


@tool
def dataframe_to_text(df: pd.DataFrame, max_rows: int = 100, max_cols: int = 100) -> str:
    """
    Convert a pandas DataFrame to a textual (string) representation.
    Shows up to max_rows and max_cols for readability.

    Args:
        df (pd.DataFrame): The DataFrame to convert.
        max_rows (int): Maximum number of rows to show (default 10).
        max_cols (int): Maximum number of columns to show (default 10).

    Returns:
        str: Textual representation of the DataFrame.
    """
    with pd.option_context('display.max_rows', max_rows, 'display.max_columns', max_cols):
        return df.to_string()


@tool
def python_bytes_to_text(python_bytes: bytes, encoding: str = 'utf-8') -> str:
    """
    Convert Python file bytes to a string representation.

    Args:
        python_bytes (bytes): The binary content of a Python file.
        encoding (str): The encoding of the Python file (default 'utf-8').

    Returns:
        str: The content of the Python file as a string.

    Raises:
        ValueError: If the bytes cannot be decoded with the specified encoding.
    """
    try:
        return python_bytes.decode(encoding)
    except Exception as e:
        raise ValueError(f"Failed to decode Python file bytes: {str(e)}")


@tool
def transcribe_audio(audio_bytes: io.BytesIO, model_size: whisperModel = 'base') -> str:
    """Transcribe audio from MP3 bytes using OpenAI's Whisper model.

    This function processes audio data and converts it to text using OpenAI's Whisper model.
    The audio is preprocessed to ensure optimal transcription quality by:
    1. Converting to mono channel
    2. Setting sample rate to 16kHz
    3. Normalizing audio samples to float32 in range [-1, 1]

    Args:
        audio_bytes (io.BytesIO): Bytes containing MP3 audio data
        model_size (whisperModel): Size of Whisper model to use. Options are:
            - 'tiny': Fastest but least accurate
            - 'base': Good balance of speed and accuracy (default)
            - 'small': More accurate but slower

    Returns:
        str: Transcribed text from the audio

    Raises:
        ValueError: If audio_bytes is invalid or corrupted
        RuntimeError: If there's an error loading the Whisper model
        Exception: For any other unexpected errors during transcription
    """
    try:
        audio = AudioSegment.from_mp3(audio_bytes)
        audio = audio.set_channels(1).set_frame_rate(16000)
        audio_array = np.array(audio.get_array_of_samples(), dtype=np.float32) / 32768.0
        model = whisper.load_model(model_size)
        result = model.transcribe(audio_array)
        return result['text']
    except Exception as e:
        if isinstance(e, ValueError):
            raise ValueError(f"Invalid audio data: {str(e)}")
        elif isinstance(e, RuntimeError):
            raise RuntimeError(f"Error loading Whisper model: {str(e)}")
        else:
            raise Exception(f"Error during transcription: {str(e)}")


@tool
def youtube_url_to_mp3(url: str) -> io.BytesIO:
    """
    Download and return MP3 audio from a given URL using yt-dlp.

    This function leverages the `yt-dlp` command-line tool to:
      1. Sanitize the provided URL by decoding percent-encoded characters and stripping invalid backslashes.
      2. Check that ffprobe is available for audio metadata extraction.
      3. Download the best available audio stream (`bestaudio`).
      4. Extract and convert the audio to MP3 format with highest quality.
      5. Load the resulting MP3 file into memory as a BytesIO object.

    Args:
        url (str): The URL of the video or audio resource to download. Can be URL-encoded.

    Returns:
        io.BytesIO: An in-memory byte-stream containing the MP3 audio data.

    Raises:
        RuntimeError: If ffprobe is missing or yt-dlp fails to download/convert the audio.
    """
    logger.debug(f"Original URL: {url}")
    clean_url = urllib.parse.unquote(url).replace("\\", "")
    logger.debug(f"Sanitized URL: {clean_url}")

    # Ensure ffprobe is available
    try:
        subprocess.run(["ffprobe", "-version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        logger.debug("ffprobe is available.")
    except Exception as e:
        logger.error(f"ffprobe check failed: {e}")
        raise RuntimeError("ffprobe is not installed or not functioning. Please install ffmpeg.")

    # Create a base temporary path without extension to let yt-dlp append .mp3
    with tempfile.NamedTemporaryFile(delete=False) as tmpf:
        base_path = tmpf.name
    os.unlink(base_path)  # remove the empty file
    output_template = base_path  # yt-dlp will add .mp3
    logger.debug(f"Using base temp path: {base_path}")

    try:
        cmd = [
            "yt-dlp",
            "-f", "bestaudio",
            "-x",
            "--audio-format", "mp3",
            "--audio-quality", "0",
            "-o", output_template,
            clean_url
        ]
        logger.debug(f"Running command: {' '.join(cmd)}")
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode != 0:
            stderr = result.stderr.strip()
            logger.error(f"yt-dlp failed: {stderr}")
            raise RuntimeError(f"yt-dlp error: {stderr}")

        mp3_path = f"{base_path}.mp3"
        logger.debug(f"Converted MP3 path: {mp3_path}")

        with open(mp3_path, 'rb') as f:
            data = f.read()
        logger.debug(f"Read {len(data)} bytes from MP3 file.")

        return io.BytesIO(data)

    finally:
        # Cleanup both files if they exist
        for path in [f"{base_path}.mp3", base_path]:
            try:
                os.remove(path)
                logger.debug(f"Removed temp file: {path}")
            except OSError:
                pass


@tool
def pdf_bytes_to_text(pdf_bytes: bytes) -> str:
    """Convert PDF bytes to text using PyPDF2.

    This tool extracts text content from a PDF file provided as bytes.
    It uses PyPDF2 to read and process the PDF document.

    Args:
        pdf_bytes (bytes): PDF file content as bytes

    Returns:
        str: Extracted text from the PDF document

    Raises:
        RuntimeError: If PDF processing fails

    Example:
        >>> pdf_bytes = open("document.pdf", "rb").read()
        >>> text = pdf_bytes_to_text(pdf_bytes)
        >>> print(text)
        "This is the text content of the PDF..."
    """
    try:
        import PyPDF2
        from io import BytesIO

        pdf_file = BytesIO(pdf_bytes)
        pdf_reader = PyPDF2.PdfReader(pdf_file)

        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text() + "\n"

        return text.strip()

    except Exception as e:
        logger.error(f"Error converting PDF to text: {e}")
        raise RuntimeError(f"Failed to convert PDF to text: {e}")