import flask
from flask import request, jsonify, send_from_directory
from flask_cors import CORS
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import json
import re
import os
import datetime

app = flask.Flask(__name__, static_folder="static")
CORS(app)

# ── Model Setup ────────────────────────────────────────────────────────────────
MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"

SEO_SYSTEM_PROMPT = """You are "Juskeo AI", an expert SEO/GEO analyst. 
Analyze the input content and return STRICTLY a raw JSON object (no markdown, no backticks) with these exact keys:
{
  "seo_title": "60-char optimized title",
  "meta_description": "155-char meta description",
  "seo_keywords": ["keyword1", "keyword2", ...],
  "focus_keyword": "primary keyword",
  "schema_type": "Article|BlogPosting|FAQPage",
  "faqs": [{"question": "...", "answer": "..."}],
  "og_title": "Open Graph title",
  "og_description": "Open Graph description",
  "canonical_hint": "suggested URL slug",
  "readability_score": "Good|Average|Poor",
  "word_count_estimate": 0,
  "geo_signals": ["city/region keywords if any"],
  "lsi_keywords": ["semantic keyword1", "semantic keyword2"]
}"""

BLOG_SYSTEM_PROMPT = """You are "Juskeo AI Content Writer". 
Write a highly detailed, comprehensive, long-form blog post. 
Include: introduction, multiple H2/H3 sections, historical context, technical details, future impacts, FAQs, conclusion.
Use markdown formatting with ##, ###, **bold**, - lists. Make it rich, 1500+ words."""

print(f"🔄 Loading {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

raw_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    trust_remote_code=True,
    low_cpu_mem_usage=True
)

print("⚡ Applying 8-bit CPU quantization...")
model = torch.quantization.quantize_dynamic(
    raw_model, {torch.nn.Linear}, dtype=torch.qint8
)
device = torch.device("cpu")
model.eval()
print("✅ Juskeo AI Engine Ready!")

# ── Helper: Run LLM ────────────────────────────────────────────────────────────
def run_llm(system_prompt, user_content, max_new_tokens=400, temperature=0.0, do_sample=False):
    chat_history = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content}
    ]
    formatted = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(formatted, return_tensors="pt", padding=True, truncation=True, max_length=2048)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature if do_sample else 1.0,
            top_p=0.9 if do_sample else 1.0,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    generated_tokens = output[0][inputs['input_ids'].shape[1]:]
    return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()

def clean_json(raw):
    """Strip markdown fences and extract JSON."""
    cleaned = re.sub(r"```(?:json)?", "", raw).replace("```", "").strip()
    # Try to find JSON object
    match = re.search(r'\{.*\}', cleaned, re.DOTALL)
    return match.group(0) if match else cleaned

def build_schema_jsonld(seo_data, article_text):
    """Build JSON-LD schema markup from SEO data."""
    schema_type = seo_data.get("schema_type", "Article")
    faqs = seo_data.get("faqs", [])
    
    schemas = []
    
    # Article/BlogPosting schema
    article_schema = {
        "@context": "https://schema.org",
        "@type": schema_type,
        "headline": seo_data.get("seo_title", ""),
        "description": seo_data.get("meta_description", ""),
        "keywords": ", ".join(seo_data.get("seo_keywords", [])),
        "datePublished": datetime.datetime.utcnow().isoformat() + "Z",
        "author": {"@type": "Organization", "name": "Juskeo AI"}
    }
    schemas.append(article_schema)
    
    # FAQ schema if FAQs exist
    if faqs:
        faq_schema = {
            "@context": "https://schema.org",
            "@type": "FAQPage",
            "mainEntity": [
                {
                    "@type": "Question",
                    "name": faq.get("question", ""),
                    "acceptedAnswer": {
                        "@type": "Answer",
                        "text": faq.get("answer", "")
                    }
                } for faq in faqs
            ]
        }
        schemas.append(faq_schema)
    
    return schemas

def generate_full_html(seo_data, blog_content, article_title):
    """Generate a complete, SEO-ready HTML page."""
    title = seo_data.get("seo_title", article_title)
    meta_desc = seo_data.get("meta_description", "")
    keywords = ", ".join(seo_data.get("seo_keywords", []))
    og_title = seo_data.get("og_title", title)
    og_desc = seo_data.get("og_description", meta_desc)
    canonical_slug = seo_data.get("canonical_hint", "article")
    schemas = build_schema_jsonld(seo_data, blog_content)
    schema_scripts = "\n".join([
        f'<script type="application/ld+json">\n{json.dumps(s, indent=2, ensure_ascii=False)}\n</script>'
        for s in schemas
    ])
    
    # Convert markdown to basic HTML
    html_body = blog_content
    html_body = re.sub(r'^### (.+)$', r'<h3>\1</h3>', html_body, flags=re.MULTILINE)
    html_body = re.sub(r'^## (.+)$', r'<h2>\1</h2>', html_body, flags=re.MULTILINE)
    html_body = re.sub(r'^# (.+)$', r'<h1>\1</h1>', html_body, flags=re.MULTILINE)
    html_body = re.sub(r'\*\*(.+?)\*\*', r'<strong>\1</strong>', html_body)
    html_body = re.sub(r'\*(.+?)\*', r'<em>\1</em>', html_body)
    html_body = re.sub(r'^- (.+)$', r'<li>\1</li>', html_body, flags=re.MULTILINE)
    html_body = re.sub(r'(<li>.*</li>\n?)+', lambda m: f'<ul>{m.group(0)}</ul>', html_body)
    html_body = re.sub(r'\n\n+', '</p><p>', html_body)
    html_body = f"<p>{html_body}</p>"
    
    faq_html = ""
    for faq in seo_data.get("faqs", []):
        faq_html += f"""
        <div class="faq-item">
            <h3 class="faq-q">{faq.get('question','')}</h3>
            <p class="faq-a">{faq.get('answer','')}</p>
        </div>"""

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>{title}</title>
    <meta name="description" content="{meta_desc}">
    <meta name="keywords" content="{keywords}">
    <meta name="robots" content="index, follow">
    <link rel="canonical" href="https://yourdomain.com/{canonical_slug}">
    
    <!-- Open Graph -->
    <meta property="og:title" content="{og_title}">
    <meta property="og:description" content="{og_desc}">
    <meta property="og:type" content="article">
    
    <!-- Twitter Card -->
    <meta name="twitter:card" content="summary_large_image">
    <meta name="twitter:title" content="{og_title}">
    <meta name="twitter:description" content="{og_desc}">
    
    <!-- Schema.org JSON-LD -->
    {schema_scripts}
    
    <style>
        body {{ font-family: Georgia, serif; max-width: 800px; margin: 0 auto; padding: 40px 20px; color: #222; line-height: 1.8; }}
        h1 {{ font-size: 2.2em; color: #111; }}
        h2 {{ font-size: 1.6em; color: #222; border-bottom: 2px solid #eee; padding-bottom: 6px; margin-top: 2em; }}
        h3 {{ font-size: 1.2em; color: #333; }}
        .faq-section {{ background: #f9f9f9; border-left: 4px solid #0066cc; padding: 20px; margin-top: 40px; }}
        .faq-q {{ color: #0066cc; margin-bottom: 4px; }}
        .faq-a {{ color: #444; margin-top: 0; }}
        ul {{ padding-left: 1.5em; }}
        li {{ margin-bottom: 6px; }}
    </style>
</head>
<body>
    <article>
        <h1>{title}</h1>
        {html_body}
        
        {'<section class="faq-section"><h2>Frequently Asked Questions</h2>' + faq_html + '</section>' if faq_html else ''}
    </article>
</body>
</html>"""


# ── AGENT PIPELINE ─────────────────────────────────────────────────────────────

@app.route('/agent/full_pipeline', methods=['POST'])
def full_pipeline():
    """
    MAIN AGENTIC ENDPOINT:
    Step 1: Extract SEO metadata from article
    Step 2: Generate/expand blog content
    Step 3: Build Schema JSON-LD
    Step 4: Assemble final SEO-ready HTML page
    """
    try:
        data = request.get_json()
        article_text = data.get("article", "").strip()
        mode = data.get("mode", "full")  # "full" | "seo_only" | "blog_only"
        
        if not article_text:
            return jsonify({"success": False, "error": "Article text required"}), 400

        pipeline_log = []
        result = {}

        # ── STEP 1: SEO Analysis ──
        pipeline_log.append("🔍 Step 1: Analyzing SEO metadata...")
        raw_seo = run_llm(SEO_SYSTEM_PROMPT, article_text[:3000], max_new_tokens=500)
        clean_seo_str = clean_json(raw_seo)
        
        try:
            seo_data = json.loads(clean_seo_str)
        except json.JSONDecodeError:
            seo_data = {"seo_title": "Optimized Article", "meta_description": "", "seo_keywords": [], "faqs": []}
        
        result["seo_metadata"] = seo_data
        pipeline_log.append("✅ SEO metadata extracted")

        # ── STEP 2: Blog Generation (if needed) ──
        blog_content = article_text  # default: use original
        if mode in ("full", "blog_only"):
            pipeline_log.append("✍️ Step 2: Expanding blog content...")
            topic_hint = seo_data.get("seo_title", article_text[:200])
            blog_content = run_llm(
                BLOG_SYSTEM_PROMPT,
                f"Topic: {topic_hint}\n\nOriginal content summary:\n{article_text[:500]}",
                max_new_tokens=1200,
                do_sample=True,
                temperature=0.7
            )
            pipeline_log.append("✅ Blog content generated")

        result["blog_content"] = blog_content

        # ── STEP 3: Schema JSON-LD ──
        pipeline_log.append("🧱 Step 3: Building Schema.org JSON-LD...")
        schemas = build_schema_jsonld(seo_data, blog_content)
        result["schema_jsonld"] = schemas
        pipeline_log.append("✅ Schema markup built")

        # ── STEP 4: Full HTML Page ──
        pipeline_log.append("🌐 Step 4: Assembling SEO-ready HTML page...")
        full_html = generate_full_html(seo_data, blog_content, seo_data.get("seo_title", "Article"))
        result["full_html"] = full_html
        pipeline_log.append("✅ HTML page assembled")

        # ── Final Response ──
        result["pipeline_log"] = pipeline_log
        result["success"] = True
        return jsonify(result)

    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


# ── Standalone endpoints (backward compat) ────────────────────────────────────

@app.route('/generate_seo', methods=['POST'])
def generate_seo():
    try:
        data = request.get_json()
        article_content = data.get("article", "")
        if not article_content:
            return jsonify({"error": "No content"}), 400
        raw_reply = run_llm(SEO_SYSTEM_PROMPT, article_content, max_new_tokens=500)
        clean_json_str = clean_json(raw_reply)
        return jsonify({"success": True, "data": clean_json_str})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/generate_blog', methods=['POST'])
def generate_blog():
    try:
        data = request.get_json()
        topic = data.get("topic", "")
        if not topic:
            return jsonify({"error": "No topic provided"}), 400
        raw_reply = run_llm(BLOG_SYSTEM_PROMPT, f"Write a massive article about: {topic}",
                            max_new_tokens=1200, do_sample=True, temperature=0.7)
        return jsonify({"success": True, "data": raw_reply})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "ok", "model": MODEL_ID, "engine": "Juskeo AI v2"})


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=7860, debug=False, threaded=True)