from dotenv import load_dotenv
load_dotenv()

from mcp.server.fastmcp import FastMCP
from transformers import pipeline, AutoTokenizer, GenerationConfig
import time

# Create the MCP server
mcp = FastMCP("ai-words", instructions="A fine-tuned language model trained on English text, dictionaries, slang, news, and reviews.")

# Load model and tokenizer once at startup
print("Loading model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
generator = pipeline(
    "text-generation",
    model="./trained_model",
    tokenizer=tokenizer,
    clean_up_tokenization_spaces=False,
)
print("Model loaded!")


@mcp.tool()
def generate_text(prompt: str, max_tokens: int = 100) -> str:
    """Generate text using the fine-tuned language model.
    
    Args:
        prompt: The text prompt to generate from.
        max_tokens: Maximum number of new tokens to generate (default 100).
    """
    config = GenerationConfig(max_new_tokens=max_tokens, do_sample=True)
    
    start_time = time.time()
    result = generator(prompt, generation_config=config)
    elapsed = time.time() - start_time
    
    generated_text = result[0]["generated_text"]
    
    input_ids = tokenizer.encode(prompt)
    output_ids = tokenizer.encode(generated_text)
    new_tokens = len(output_ids) - len(input_ids)
    speed = new_tokens / elapsed if elapsed > 0 else 0
    
    return (
        f"{generated_text}\n\n"
        f"--- Stats ---\n"
        f"Prompt tokens: {len(input_ids)}\n"
        f"New tokens: {new_tokens}\n"
        f"Time: {elapsed:.2f}s\n"
        f"Speed: {speed:.1f} tokens/sec"
    )


@mcp.tool()
def define_word(word: str) -> str:
    """Look up a word or slang definition using the trained model.
    
    Args:
        word: The word or phrase to define.
    """
    config = GenerationConfig(max_new_tokens=150, do_sample=True)
    result = generator(f"{word}:", generation_config=config)
    return result[0]["generated_text"]


@mcp.tool()
def complete_sentence(partial: str) -> str:
    """Complete a partial sentence or paragraph.
    
    Args:
        partial: The beginning of a sentence or paragraph to complete.
    """
    config = GenerationConfig(max_new_tokens=200, do_sample=True)
    result = generator(partial, generation_config=config)
    return result[0]["generated_text"]


@mcp.tool()
def model_info() -> str:
    """Get information about the trained model."""
    vocab_size = len(tokenizer)
    return (
        f"Model: GPT-2 fine-tuned with Qwen3-8B tokenizer\n"
        f"Tokenizer: Qwen/Qwen3-8B\n"
        f"Vocabulary size: {vocab_size:,}\n"
        f"Model path: ./trained_model\n"
        f"Training data: 13 datasets (~2.9M examples)\n"
        f"  - WikiText-103, WikiText-2\n"
        f"  - English Dictionary, WordNet\n"
        f"  - AG News, IMDB, Rotten Tomatoes\n"
        f"  - CNN/DailyMail, Yelp Reviews\n"
        f"  - Urban Dictionary, Slang, Gen Z Slang"
    )


if __name__ == "__main__":
    mcp.run()