import openai
import gradio as gr
from typing import Dict, List
import re
from humanize import paraphrase_text
from ai_generate import generate
import requests
from gptzero_free import GPT2PPL


def clean_text(text: str) -> str:
    paragraphs = text.split('\n\n')
    cleaned_paragraphs = []
    for paragraph in paragraphs:
        cleaned = re.sub(r'\s+', ' ', paragraph).strip()
        cleaned = re.sub(r'(?<=\.) ([a-z])', lambda x: x.group(1).upper(), cleaned)
        cleaned_paragraphs.append(cleaned)
    return '\n'.join(cleaned_paragraphs)

def format_and_correct(text: str) -> str:
    """Correct formatting and grammar without changing content significantly."""
    prompt = f"""
    Please correct the formatting, grammar, and spelling errors in the following text without changing its content significantly. Ensure proper paragraph breaks and maintain the original content:
    {text}
    """
    corrected_text = generate(prompt, "Groq", None) 
    return clean_text(corrected_text)


def generate_prompt(settings: Dict[str, str]) -> str:
    """Generate a detailed prompt based on user settings."""
    prompt = f"""
    Write a {settings['article_length']} {settings['format']} on {settings['topic']}.
    
    Style and Tone:
    - Writing style: {settings['writing_style']}
    - Tone: {settings['tone']}
    - Target audience: {settings['user_category']}
    
    Content:
    - Depth: {settings['depth_of_content']}
    - Structure: {', '.join(settings['structure'])}
    
    Keywords to incorporate:
    {', '.join(settings['keywords'])}
    
    Additional requirements:
    - Include {settings['num_examples']} relevant examples or case studies
    - Incorporate data or statistics from {', '.join(settings['references'])}
    - End with a {settings['conclusion_type']} conclusion
    - Add a "References" section at the end with at least 3 credible sources, formatted as [1], [2], etc.
    - Do not make any headline, title bold.
    
    Ensure proper paragraph breaks for better readability.
    Avoid any references to artificial intelligence, language models, or the fact that this is generated by an AI, and do not mention something like here is the article etc.
    """
    return prompt

def generate_article(
    topic: str,
    keywords: str,
    article_length: str,
    format: str,
    writing_style: str,
    tone: str,
    user_category: str,
    depth_of_content: str,
    structure: str,
    references: str,
    num_examples: str,
    conclusion_type: str,
    ai_model: str,
    api_key: str = None
) -> str:
    """Generate an article based on user-defined settings."""
    settings = {
        "topic": topic,
        "keywords": [k.strip() for k in keywords.split(',')],
        "article_length": article_length,
        "format": format,
        "writing_style": writing_style,
        "tone": tone,
        "user_category": user_category,
        "depth_of_content": depth_of_content,
        "structure": [s.strip() for s in structure.split(',')],
        "references": [r.strip() for r in references.split(',')],
        "num_examples": num_examples,
        "conclusion_type": conclusion_type
    }
    
    prompt = generate_prompt(settings)
    
    if ai_model in ['OpenAI GPT 3.5', 'OpenAI GPT 4']:
        response = openai.ChatCompletion.create(
            model="gpt-4" if ai_model == 'OpenAI GPT 4' else "gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a professional content writer with expertise in various fields."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=3000,
            n=1,
            stop=None,
            temperature=0.7,
        )
        article = response.choices[0].message.content.strip()
    else:
        article = generate(prompt, ai_model, api_key)
    
    return clean_text(article)

def humanize(
    text: str,    
    model: str, 
    temperature: float = 1.2, 
    repetition_penalty: float = 1,
    top_k: int = 50,
    length_penalty: float = 1
) -> str:
    result = paraphrase_text(
        text=text,
        model_name=model,
        temperature=temperature,
        repetition_penalty=repetition_penalty,
        top_k=top_k,
        length_penalty=length_penalty,
    )
    return format_and_correct(result)

ai_check_options = [
    "Polygraf AI",
    # "Sapling AI",
    "GPTZero"
]

def ai_generated_test_polygraf(text: str) -> Dict:
    url = "http://34.66.10.188/ai-vs-human"
    access_key = "6mcemwsFycVVgVjMFwKXki3zJka1r7N4u$Z0Y|x$gecC$hdNtpQf-SpL0+=k;u%BZ" 
    headers = {
        "ACCESS_KEY": access_key
    }
    data = {
    "text" : f"{text}"
    }
    response = requests.post(url, headers=headers, json=data)
    return response.json()

def ai_generated_test_sapling(text: str) -> Dict:
    response = requests.post(
        "https://api.sapling.ai/api/v1/aidetect",
        json={
            "key": "60L9BPSVPIIOEZM0CD1DQWRBPJIUR7SB",
            "text": f"{text}"
        }
    )
    return { "AI" : response.json()['score'], "HUMAN" : 1 - response.json()['score']}


def ai_generated_test_gptzero(text):
    gptzero_model = GPT2PPL()
    result = gptzero_model(text)
    print(result)
    return result

def ai_check(text: str, option: str) -> Dict:
    if option == 'Polygraf AI':
        return ai_generated_test_polygraf(text)
    elif option == 'Sapling AI':
        return ai_generated_test_sapling(text)
    elif option == "GPTZero":
        return ai_generated_test_gptzero(text)
    else:
        return ai_generated_test_polygraf(text)

def update_visibility_api(model: str):
    if model in ['OpenAI GPT 3.5', 'OpenAI GPT 4']:
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)

def format_references(text: str) -> str:
    """Extract and format references from the generated text."""
    lines = text.split('\n')
    references = []
    article_text = []
    in_references = False
    
    for line in lines:
        if line.strip().lower() == "references":
            in_references = True
            continue
        if in_references:
            references.append(line.strip())
        else:
            article_text.append(line)
    
    formatted_refs = []
    for i, ref in enumerate(references, 1):
        formatted_refs.append(f"[{i}] {ref}\n")
        
    return "\n\n".join(article_text) + "\n\nReferences:\n" + "\n".join(formatted_refs)

def generate_and_format(
    topic, keywords, article_length, format, writing_style, tone, user_category,
    depth_of_content, structure, references, num_examples, conclusion_type, ai_model, api_key
):
    article = generate_article(
        topic, keywords, article_length, format, writing_style, tone, user_category,
        depth_of_content, structure, references, num_examples, conclusion_type, ai_model, api_key
    )
    return format_references(article)

def copy_to_input(text):
    return text

def create_interface():
    with gr.Blocks(theme=gr.themes.Default(
        primary_hue=gr.themes.colors.pink,
        secondary_hue=gr.themes.colors.yellow,
        neutral_hue=gr.themes.colors.gray
    )) as demo:
        gr.Markdown("# Polygraf AI Content Writer", elem_classes="text-center text-3xl mb-6")
        
        with gr.Row():
            with gr.Column(scale=2):
                with gr.Group():
                    gr.Markdown("## Article Configuration", elem_classes="text-xl mb-4")
                    input_topic = gr.Textbox(label="Topic", placeholder="Enter the main topic of your article", elem_classes="input-highlight-pink")
                    input_keywords = gr.Textbox(label="Keywords", placeholder="Enter comma-separated keywords", elem_classes="input-highlight-yellow")
                    
                    with gr.Row():
                        input_format = gr.Dropdown(
                            choices=['Article', 'Essay', 'Blog post', 'Report', 'Research paper', 'News article', 'White paper'],
                            value='Article',
                            label="Format",
                            elem_classes="input-highlight-turquoise"
                        )
                        input_length = gr.Dropdown(
                            choices=["Short (500 words)", "Medium (1000 words)", "Long (2000+ words)", "Very Long (3000+ words)"],
                            value="Medium (1000 words)",
                            label="Article Length",
                            elem_classes="input-highlight-pink"
                        )
                    
                    with gr.Row():
                        input_writing_style = gr.Dropdown(
                            choices=["Formal", "Informal", "Technical", "Conversational", "Journalistic", "Academic", "Creative"],
                            value="Formal",
                            label="Writing Style",
                            elem_classes="input-highlight-yellow"
                        )
                        input_tone = gr.Dropdown(
                            choices=["Friendly", "Professional", "Neutral", "Enthusiastic", "Skeptical", "Humorous"],
                            value="Professional",
                            label="Tone",
                            elem_classes="input-highlight-turquoise"
                        )
                    
                    input_user_category = gr.Dropdown(
                        choices=["Students", "Professionals", "Researchers", "General Public", "Policymakers", "Entrepreneurs"],
                        value="General Public",
                        label="Target Audience",
                        elem_classes="input-highlight-pink"
                    )
                    input_depth = gr.Dropdown(
                        choices=["Surface-level overview", "Moderate analysis", "In-depth research", "Comprehensive study"],
                        value="Moderate analysis",
                        label="Depth of Content",
                        elem_classes="input-highlight-yellow"
                    )
                    input_structure = gr.Dropdown(
                        choices=[
                            "Introduction, Body, Conclusion",
                            "Abstract, Introduction, Methods, Results, Discussion, Conclusion",
                            "Executive Summary, Problem Statement, Analysis, Recommendations, Conclusion",
                            "Introduction, Literature Review, Methodology, Findings, Analysis, Conclusion"
                        ],
                        value="Introduction, Body, Conclusion",
                        label="Structure",
                        elem_classes="input-highlight-turquoise"
                    )
                    input_references = gr.Dropdown(
                        choices=["Academic journals", "Industry reports", "Government publications", "News outlets", "Expert interviews", "Case studies"],
                        value="News outlets",
                        label="References",
                        elem_classes="input-highlight-pink"
                    )
                    input_num_examples = gr.Dropdown(
                        choices=["1-2", "3-4", "5+"],
                        value="1-2",
                        label="Number of Examples/Case Studies",
                        elem_classes="input-highlight-yellow"
                    )
                    input_conclusion = gr.Dropdown(
                        choices=["Summary", "Call to Action", "Future Outlook", "Thought-provoking Question"],
                        value="Summary",
                        label="Conclusion Type",
                        elem_classes="input-highlight-turquoise"
                    )
                
                with gr.Group():
                    gr.Markdown("## AI Model Configuration", elem_classes="text-xl mb-4")
                    ai_generator = gr.Dropdown(
                        choices=['Llama 3', 'Groq', 'Mistral', 'Gemma', 'OpenAI GPT 3.5', 'OpenAI GPT 4'],
                        value='Llama 3',
                        label="AI Model",
                        elem_classes="input-highlight-pink"
                    )
                    input_api = gr.Textbox(label="API Key", visible=False)
                    ai_generator.change(update_visibility_api, ai_generator, input_api)
                
                generate_btn = gr.Button("Generate Article", variant="primary")
            
            with gr.Column(scale=3):
                output_article = gr.Textbox(label="Generated Article", lines=20)
                
                with gr.Row():
                    with gr.Column():
                        ai_detector_dropdown = gr.Radio(
                            choices=ai_check_options, label="Select AI Detector", value="Polygraf AI")
                        ai_check_btn = gr.Button("AI Check")
                    ai_check_result = gr.Label(label="AI Check Result")
                
                humanize_btn = gr.Button("Humanize")
                humanized_output = gr.Textbox(label="Humanized Article", lines=20)
                copy_to_input_btn = gr.Button("Copy to Input for AI Check")
        
        with gr.Accordion("Advanced Humanizer Settings", open=False):
            with gr.Row():
                model_dropdown = gr.Radio(
                    choices=[
                        "Base Model",
                        "Large Model",
                        "XL Model",
                        # "XL Law Model",
                        # "XL Marketing Model",
                        # "XL Child Style Model",
                    ], 
                    value="Large Model",
                    label="Humanizer Model Version"
                )
            with gr.Row():
                temperature_slider = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.2, label="Temperature")
                top_k_slider = gr.Slider(minimum=0, maximum=300, step=25, value=50, label="Top k")
            with gr.Row():
                repetition_penalty_slider = gr.Slider(minimum=1.0, maximum=2.0, step=0.1, value=1, label="Repetition Penalty")
                length_penalty_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1.0, label="Length Penalty")

        generate_btn.click(
            fn=generate_and_format,
            inputs=[
                input_topic, 
                input_keywords,
                input_length,
                input_format,
                input_writing_style,
                input_tone,
                input_user_category,
                input_depth,
                input_structure,
                input_references,
                input_num_examples,
                input_conclusion,
                ai_generator,
                input_api
            ],
            outputs=[output_article],
        )

        ai_check_btn.click(
            fn=ai_check,
            inputs=[output_article, ai_detector_dropdown],
            outputs=[ai_check_result],
        )

        humanize_btn.click(
            fn=humanize,
            inputs=[
                output_article,
                model_dropdown,
                temperature_slider,
                repetition_penalty_slider,
                top_k_slider,
                length_penalty_slider,
            ],
            outputs=[humanized_output],
        )

        copy_to_input_btn.click(
            fn=copy_to_input,
            inputs=[humanized_output],
            outputs=[output_article],
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(server_name="0.0.0.0", share=True)