import subprocess
import sys
from transformers import T5ForConditionalGeneration, T5Tokenizer
import gradio as gr

# Install required packages
def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install('transformers')
install('sentencepiece')
install('torch')
install('gradio')

# Translation setup
try:
    install('translate')
    from translate import Translator
    translation_available = True
except:
    print("Translation package not available - continuing without translation support")
    translation_available = False

class GOTSummarizer:
    def __init__(self):
        self.tokenizer = T5Tokenizer.from_pretrained("t5-small")
        self.model = T5ForConditionalGeneration.from_pretrained("t5-small")
        self.translation_available = translation_available
    
    def summarize(self, text, max_length=150):
        inputs = self.tokenizer("summarize: " + text, 
                              return_tensors="pt", 
                              truncation=True, 
                              max_length=512)
        outputs = self.model.generate(
            **inputs,
            max_length=max_length,
            num_beams=4,
            early_stopping=True
        )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    def translate(self, text, lang='hi'):
        if not self.translation_available:
            return "[Translation unavailable] " + text
            
        try:
            translator = Translator(to_lang=lang)  # Initialize with target language
            return translator.translate(text)
        except Exception as e:
            print(f"Translation error: {e}")
            return text

# Initialize summarizer
summarizer = GOTSummarizer()

def process(text, lang='en'):
    if not text.strip():
        return "Please enter text to summarize"
    
    summary = summarizer.summarize(text)
    return summary if lang == 'en' else summarizer.translate(summary, lang)

# Create Gradio interface
interface = gr.Interface(
    fn=process,
    inputs=[
        gr.Textbox(label="Game of Thrones Text", lines=10,
                 placeholder="Paste book chapter text here..."),
        gr.Dropdown(
            label="Language",
            choices=['en', 'hi', 'pa', 'ta', 'bn'],
            value='en'
        )
    ],
    outputs=gr.Textbox(label="Summary"),
    title="ASOIAF Chapter Summarizer",
    description="Summarizes Game of Thrones chapters with optional translation"
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()