import logging
import textwrap
import threading
from typing import Literal, Optional, Tuple, Union

import gradio as gr
import outlines
import pandas as pd
import spaces
import torch
from outlines import generate
from peft import PeftConfig, PeftModel
from pydantic import BaseModel, ConfigDict
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration
MODEL_CACHE = {}
MODEL_LOCK = threading.Lock()
DEVICE_MAP = "auto"
QUANTIZATION_BITS = 4  # Changed to 4-bit by default for efficiency
TEMPERATURE = 0.0

AVAILABLE_MODELS = [
    "rshwndsz/ft-longformer-base-4096",
    "rshwndsz/ft-hermes-3-llama-3.2-3b",
    "rshwndsz/ft-phi-3.5-mini-instruct",
    "rshwndsz/ft-mistral-7b-v0.3-instruct",
    "rshwndsz/ft-phi-4",
    "rshwndsz/ft_paraphrased-hermes-3-llama-3.2-3b",
    "rshwndsz/ft_paraphrased-longformer-base-4096",
    "rshwndsz/ft_paraphrased-phi-3.5-mini-instruct",
    "rshwndsz/ft_paraphrased-mistral-7b-v0.3-instruct",
    "rshwndsz/ft_paraphrased-phi-4",
]
DEFAULT_MODEL_ID = AVAILABLE_MODELS[0]

SYSTEM_PROMPT = textwrap.dedent("""
You are an assistant tasked with grading answers to a mind reading ability test. You will be provided with the following information:
1. A story that was presented to participants as context
2. The question that participants were asked to answer
3. A grading scheme to evaluate the answers (Correct Responses:1, incorrect response:0, Incomplete response:0, Irrelevant:0)
4. Grading examples
5. A participant answer
Your task is to grade each answer according to the grading scheme. For each answer, you should:
1. Carefully read and understand the answer and compare it to the grading criteria
2. Assigning an score 1 or 0 for each answer.
""").strip()

PROMPT_TEMPLATE = textwrap.dedent("""
<Story>
{story}
</Story>
<Question>
{question}
</Question>
<GradingScheme>
{grading_scheme}
</GradingScheme>
<Answer>
{answer}
</Answer>
Score:""").strip()

class ResponseModel(BaseModel):
    model_config = ConfigDict(extra="forbid")
    score: Literal["0", "1"]

def get_model_and_tokenizer(
    model_id: str, 
    device_map: str = "auto", 
    quantization_bits: Optional[int] = 4
) -> Tuple[Union[AutoModelForCausalLM, AutoModelForSequenceClassification], AutoTokenizer]:
    """Load model and tokenizer with caching"""
    with MODEL_LOCK:
        if model_id in MODEL_CACHE:
            return MODEL_CACHE[model_id]
        
        try:
            if quantization_bits == 4:
                quantization_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_use_double_quant=True,
                    bnb_4bit_compute_dtype=torch.bfloat16,
                )
            elif quantization_bits == 8:
                quantization_config = BitsAndBytesConfig(load_in_8bit=True)
            else:
                quantization_config = None

            if "longformer" in model_id:
                # For sequence classification models
                model = AutoModelForSequenceClassification.from_pretrained(
                    model_id, 
                    device_map=device_map
                )
                tokenizer = AutoTokenizer.from_pretrained(model_id)
                if tokenizer.pad_token is None:
                    tokenizer.pad_token = tokenizer.eos_token
            else:
                # For causal LM models
                peft_config = PeftConfig.from_pretrained(model_id)
                base_model_id = peft_config.base_model_name_or_path
                
                model = AutoModelForCausalLM.from_pretrained(
                    base_model_id,
                    device_map=device_map,
                    quantization_config=quantization_config,
                    torch_dtype=torch.bfloat16,
                )
                model = PeftModel.from_pretrained(model, model_id)
                tokenizer = AutoTokenizer.from_pretrained(
                    base_model_id, 
                    use_fast=True, 
                    clean_up_tokenization_spaces=True
                )
                if tokenizer.pad_token is None:
                    tokenizer.pad_token = tokenizer.eos_token

            MODEL_CACHE[model_id] = (model, tokenizer)
            return model, tokenizer
            
        except Exception as e:
            logger.error(f"Error loading model {model_id}: {str(e)}")
            raise

def format_prompt(story: str, question: str, grading_scheme: str, answer: str) -> str:
    prompt = PROMPT_TEMPLATE.format(
        story=story.strip(),
        question=question.strip(),
        grading_scheme=grading_scheme.strip(),
        answer=answer.strip(),
    )
    full_prompt = SYSTEM_PROMPT + "\n\n" + prompt
    return full_prompt

@spaces.GPU
def label_single_response_with_model(model_id, story, question, criteria, response):
    try:
        prompt = format_prompt(story, question, criteria, response)
        model, tokenizer = get_model_and_tokenizer(model_id, DEVICE_MAP, QUANTIZATION_BITS)

        if "longformer" in model_id:
            # Sequence classification approach
            inputs = tokenizer(
                prompt, 
                return_tensors="pt", 
                truncation=True, 
                padding=True,
                max_length=4096
            )
            with torch.no_grad():
                logits = model(**inputs).logits
            predicted_class = torch.argmax(logits, dim=1).item()
            return str(predicted_class)
        else:
            # Structured generation with outlines
            generator = generate.json(model, ResponseModel, max_tokens=20)
            result = generator(prompt)
            return result.score
            
    except Exception as e:
        logger.error(f"Error in single response labeling: {str(e)}")
        return f"Error: {str(e)}"

@spaces.GPU
def label_multi_responses_with_model(model_id, story, question, criteria, response_file):
    try:
        df = pd.read_csv(response_file.name)
        assert "response" in df.columns, "CSV must contain a 'response' column."
        
        model, tokenizer = get_model_and_tokenizer(model_id, DEVICE_MAP, QUANTIZATION_BITS)
        scores = []
        
        if "longformer" in model_id:
            # Batch processing for sequence classification
            prompts = [
                format_prompt(story, question, criteria, resp) 
                for resp in df["response"]
            ]
            inputs = tokenizer(
                prompts, 
                return_tensors="pt", 
                truncation=True, 
                padding=True,
                max_length=4096
            )
            with torch.no_grad():
                logits = model(**inputs).logits
            predicted_classes = torch.argmax(logits, dim=1).tolist()
            scores = [str(cls) for cls in predicted_classes]
        else:
            # Sequential processing for generative models
            generator = generate.json(model, ResponseModel, max_tokens=20)
            for response in df["response"]:
                prompt = format_prompt(story, question, criteria, response)
                result = generator(prompt)
                scores.append(result.score)
        
        df["score"] = scores
        return df
        
    except Exception as e:
        logger.error(f"Error in multi response labeling: {str(e)}")
        return pd.DataFrame({"error": [str(e)]})

def single_response_ui(model_id):
    return gr.Interface(
        fn=lambda story, question, criteria, response: label_single_response_with_model(
            model_id, story, question, criteria, response
        ),
        inputs=[
            gr.Textbox(label="Story", lines=6),
            gr.Textbox(label="Question", lines=2),
            gr.Textbox(label="Criteria (Grading Scheme)", lines=4),
            gr.Textbox(label="Single Response", lines=3),
        ],
        outputs=gr.Textbox(label="Score"),
        live=False,
        title="Single Response Grader",
        description="Grade a single response against the story, question, and criteria"
    )

def multi_response_ui(model_id):
    return gr.Interface(
        fn=lambda story, question, criteria, response_file: label_multi_responses_with_model(
            model_id, story, question, criteria, response_file
        ),
        inputs=[
            gr.Textbox(label="Story", lines=6),
            gr.Textbox(label="Question", lines=2),
            gr.Textbox(label="Criteria (Grading Scheme)", lines=4),
            gr.File(
                label="Responses CSV (.csv with 'response' column)", 
                file_types=[".csv"]
            ),
        ],
        outputs=gr.Dataframe(label="Labeled Responses", type="pandas"),
        live=False,
        title="Batch Response Grader",
        description="Upload a CSV file with responses to grade them in batch"
    )

with gr.Blocks(title="Zero-Shot Evaluation Grader") as iface:
    gr.Markdown("# Zero-Shot Evaluation Grader")
    gr.Markdown("Select a model and then use either the single response or batch processing tab.")
    
    model_selector = gr.Dropdown(
        label="Select Model",
        choices=AVAILABLE_MODELS,
        value=DEFAULT_MODEL_ID,
    )
    
    with gr.Tabs():
        with gr.Tab("Single Response"):
            single_response_ui(model_selector.value)
        with gr.Tab("Batch Processing (CSV)"):
            multi_response_ui(model_selector.value)

if __name__ == "__main__":
    iface.launch(share=True)