import logging
import torch
import base64
import os

from indexify_extractor_sdk import Content, Extractor, Feature
from pyannote.audio import Pipeline
from transformers import pipeline, AutoModelForCausalLM
from .diarization_utils import diarize
from huggingface_hub import HfApi
from starlette.exceptions import HTTPException

from pydantic import BaseModel
from pydantic_settings import BaseSettings
from typing import Optional, Literal, List, Union

logger = logging.getLogger(__name__)
token = os.getenv('HF_TOKEN')

class ModelSettings(BaseSettings):
    asr_model: str = "openai/whisper-large-v3"
    assistant_model: Optional[str] = "distil-whisper/distil-large-v3"
    diarization_model: Optional[str] = "pyannote/speaker-diarization-3.1"
    hf_token: Optional[str] = token

model_settings = ModelSettings()

class ASRExtractorConfig(BaseModel):
    task: Literal["transcribe", "translate"] = "transcribe"
    batch_size: int = 24
    assisted: bool = False
    chunk_length_s: int = 30
    sampling_rate: int = 16000
    language: Optional[str] = None
    num_speakers: Optional[int] = None
    min_speakers: Optional[int] = None
    max_speakers: Optional[int] = None

class ASRExtractor(Extractor):
    name = "tensorlake/asrdiarization"
    description = "Powerful ASR + diarization + speculative decoding."
    system_dependencies = ["ffmpeg"]
    input_mime_types = ["audio", "audio/mpeg"]

    def __init__(self):
        super(ASRExtractor, self).__init__()

        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        logger.info(f"Using device: {device.type}")
        torch_dtype = torch.float32 if device.type == "cpu" else torch.float16

        self.assistant_model = AutoModelForCausalLM.from_pretrained(
            model_settings.assistant_model,
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True,
            use_safetensors=True
        ) if model_settings.assistant_model else None

        if self.assistant_model:
            self.assistant_model.to(device)

        self.asr_pipeline = pipeline(
            "automatic-speech-recognition",
            model=model_settings.asr_model,
            torch_dtype=torch_dtype,
            device=device
        )

        if model_settings.diarization_model:
            # diarization pipeline doesn't raise if there is no token
            HfApi().whoami(model_settings.hf_token)
            self.diarization_pipeline = Pipeline.from_pretrained(
                checkpoint_path=model_settings.diarization_model,
                use_auth_token=model_settings.hf_token,
            )
            self.diarization_pipeline.to(device)
        else:
            self.diarization_pipeline = None

    def extract(self, content: Content, params: ASRExtractorConfig) -> List[Union[Feature, Content]]:
        file = base64.b64decode(content.data)
        logger.info(f"inference params: {params}")

        generate_kwargs = {
            "task": params.task, 
            "language": params.language,
            "assistant_model": self.assistant_model if params.assisted else None
        }

        try:
            asr_outputs = self.asr_pipeline(
                file,
                chunk_length_s=params.chunk_length_s,
                batch_size=params.batch_size,
                generate_kwargs=generate_kwargs,
                return_timestamps=True,
            )
        except RuntimeError as e:
            logger.error(f"ASR inference error: {str(e)}")
            raise HTTPException(status_code=400, detail=f"ASR inference error: {str(e)}")
        except Exception as e:
            logger.error(f"Unknown error diring ASR inference: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Unknown error diring ASR inference: {str(e)}")

        if self.diarization_pipeline:
            try:
                transcript = diarize(self.diarization_pipeline, file, params, asr_outputs)
            except RuntimeError as e:
                logger.error(f"Diarization inference error: {str(e)}")
                raise HTTPException(status_code=400, detail=f"Diarization inference error: {str(e)}")
            except Exception as e:
                logger.error(f"Unknown error during diarization: {str(e)}")
                raise HTTPException(status_code=500, detail=f"Unknown error during diarization: {str(e)}")
        else:
            transcript = []

        feature = Feature.metadata(value={"chunks": asr_outputs["chunks"], "text": asr_outputs["text"]})
        return [Content.from_text(str(transcript), features=[feature])]
    
    def sample_input(self) -> Content:
        filepath = "sample.mp3"
        with open(filepath, 'rb') as f:
            audio_encoded = base64.b64encode(f.read()).decode("utf-8")
        return Content(content_type="audio/mpeg", data=audio_encoded)

if __name__ == "__main__":
    filepath = "sample.mp3"
    with open(filepath, 'rb') as f:
        audio_encoded = base64.b64encode(f.read()).decode("utf-8")
    data = Content(content_type="audio/mpeg", data=audio_encoded)
    params = ASRExtractorConfig(batch_size=24)
    extractor = ASRExtractor()
    results = extractor.extract(data, params=params)
    print(results)