import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import logging
import sys
import json
from datetime import datetime

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

logger.info(f"===== Application Startup at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} =====")

# Initialize model and tokenizer
logger.info("Loading model and tokenizer...")
model_name = "deepseek-ai/deepseek-coder-1.3b-base"  # Using smaller model for CPU

try:
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True,
        cache_dir="/tmp/deepseek_cache"
    )
    logger.info("Tokenizer loaded successfully")
    
    # Load model in CPU mode
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32,  # Use float32 for CPU
        device_map=None,  # Don't use device map
        trust_remote_code=True,
        cache_dir="/tmp/deepseek_cache"
    ).to('cpu')  # Explicitly move to CPU
    
    logger.info("Model loaded successfully")
    logger.info(f"Model device: {next(model.parameters()).device}")
    logger.info(f"Total model parameters: {sum(p.numel() for p in model.parameters())}")
except Exception as e:
    logger.error(f"Error loading model: {str(e)}")
    logger.error("Traceback:", exc_info=True)
    raise

def generate_tweet(style_profile, recent_tweets, topic):
    logger.info(f"Generating tweet for topic: {topic}")
    logger.info(f"Style Profile: {json.dumps(style_profile)}")
    logger.info(f"Recent Tweets: {recent_tweets}")
    
    try:
        # Construct the prompt
        prompt = f"""Generate a tweet about "{topic}" that matches the following style:

Style Characteristics:
Tone: {style_profile.get('tone', {}).get('dominantTone', 'Natural')}
Complexity: {style_profile.get('complexity', 2.8)}
Technical Level: {style_profile.get('summary', {}).get('technicalLevel', 'Technical')}

Recent tweet examples for style reference:
{recent_tweets}

Generate a single tweet (max 280 characters) that matches this style and discusses the given topic.
Response format: Just the tweet text, no additional commentary."""

        logger.info(f"Generated prompt: {prompt[:200]}...")

        # Generate the tweet
        inputs = tokenizer(prompt, return_tensors="pt")
        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                max_new_tokens=280,
                temperature=0.7,
                top_p=0.95,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        generated_tweet = tokenizer.decode(outputs[0], skip_special_tokens=True)
        logger.info(f"Generated tweet: {generated_tweet}")
        
        return {
            "content": generated_tweet,
            "metadata": {
                "model": model_name,
                "prompt_length": len(prompt),
                "response_length": len(generated_tweet)
            }
        }
    except Exception as e:
        logger.error(f"Error in generate_tweet: {str(e)}")
        logger.error("Traceback:", exc_info=True)
        raise

def process_request(style_profile_str, recent_tweets_str, topic):
    """Process the request with string inputs and convert them to the required format"""
    logger.info("Processing request...")
    try:
        # Parse the style profile from string to dict if needed
        if isinstance(style_profile_str, str):
            logger.info("Parsing style profile from string")
            style_profile = json.loads(style_profile_str)
        else:
            style_profile = style_profile_str
            
        # Split recent tweets if it's a string
        if isinstance(recent_tweets_str, str):
            recent_tweets = recent_tweets_str.strip()
        else:
            recent_tweets = "\n".join(tweet.get('text', '') for tweet in recent_tweets_str)
            
        logger.info("Calling generate_tweet...")
        result = generate_tweet(style_profile, recent_tweets, topic)
        logger.info("Successfully generated tweet")
        return result
    except Exception as e:
        logger.error(f"Error in process_request: {str(e)}")
        logger.error("Traceback:", exc_info=True)
        raise

# Create Gradio interface
logger.info("Creating Gradio interface...")

with gr.Blocks(title="DeepSeek Tweet Generator") as iface:
    gr.Markdown("# DeepSeek Tweet Generator")
    gr.Markdown("Generate tweets using DeepSeek model based on style profile and recent tweets")
    
    with gr.Row():
        with gr.Column():
            style_profile = gr.JSON(
                label="Style Profile",
                value={"tone": {"dominantTone": "Technical"}, "complexity": 2.8, "summary": {"technicalLevel": "Technical"}}
            )
            recent_tweets = gr.Textbox(
                label="Recent Tweets",
                lines=5,
                value="Just deployed a new microservices architecture using Kubernetes!\nExcited to share my latest ML model trained on AWS SageMaker."
            )
            topic = gr.Textbox(
                label="Topic",
                value="Web3 Development"
            )
            generate_btn = gr.Button("Generate Tweet")
        
        with gr.Column():
            output = gr.JSON(label="Generated Tweet")
    
    generate_btn.click(
        fn=process_request,
        inputs=[style_profile, recent_tweets, topic],
        outputs=output
    )

# Launch the interface
logger.info("Launching Gradio interface...")
iface.launch()