"""
Content processing and rewriting using configurable AI providers.
"""
import json
import requests
import tiktoken
from types import SimpleNamespace
from groq import Groq

import json
import requests
from types import SimpleNamespace


import os
from huggingface_hub import InferenceClient

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None
from logger import logger
from config import (
    AI_PROVIDER,
    AI_MODEL,
    CONTENT_TEMPERATURE,
    MAX_CONTENT_PREVIEW,
    PROMPTS_DIR,
    GROQ_API_KEY,
    GOOGLE_API_KEY,
    DEEPSEEK_API_KEY,
    OPENAI_API_KEY,
    HUGGINGFACE_API_KEY,
    HUGGINGFACE_ENDPOINT,
    HF_INFERENCE_PROVIDER,
)


def _load_prompt_template(filename):
    """Load a prompt template from the prompts directory."""
    path = os.path.join(PROMPTS_DIR, filename)
    with open(path, "r", encoding="utf-8") as f:
        return f.read()

class GoogleAIClient:
    """Lightweight wrapper for Google Generative AI REST API."""

    def __init__(self, api_key):
        self.api_key = api_key
        self.chat = self
        self.completions = self

    def create(self, model, messages, temperature, response_format=None):
        prompt_text = "\n".join(
            f"{message['role']}: {message['content']}" for message in messages
        )
        url = (
            f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={self.api_key}"
        )
        payload = {
            "contents": [
                {
                    "parts": [
                        {"text": prompt_text}
                    ]
                }
            ],
            "generationConfig": {
                "temperature": temperature,
            },
        }

        response = requests.post(
            url,
            headers={"Content-Type": "application/json"},
            json=payload,
            timeout=30,
        )
        response.raise_for_status()

        data = response.json()
        content = self._extract_response_text(data)

        return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content=content))])

    def _extract_response_text(self, data):
        if not data:
            return ""

        # New Google responses may return candidates or output content structures.
        candidates = data.get("candidates") or data.get("candidates", [])
        if candidates:
            candidate = candidates[0]
            content = candidate.get("content") or candidate.get("output") or candidate.get("text")
            return self._normalize_content(content)

        output = data.get("output")
        if output is not None:
            return self._normalize_content(output)

        return ""

    def _normalize_content(self, content):
        if content is None:
            return ""
        if isinstance(content, str):
            return content
        if isinstance(content, dict):
            if "text" in content:
                return str(content["text"])
            if "output" in content:
                return self._normalize_content(content["output"])
            if "contents" in content:
                return self._normalize_content(content["contents"])
            if "candidates" in content:
                return self._normalize_content(content["candidates"])
            if "parts" in content:
                return self._normalize_content(content["parts"])
            # Flatten dict values
            for value in content.values():
                normalized = self._normalize_content(value)
                if normalized:
                    return normalized
            return ""
        if isinstance(content, list):
            parts = [self._normalize_content(item) for item in content]
            return " ".join([part for part in parts if part])
        return str(content)


class ContentProcessor:
    """Handles content rewriting and tag generation using configurable AI providers."""

    def __init__(
        self,
        provider=AI_PROVIDER,
        model=AI_MODEL,
        temperature=CONTENT_TEMPERATURE,
        api_key=None,
    ):
        """Initialize AI client and model settings."""
        self.provider = provider.lower()
        self.model = model
        self.temperature = temperature
        self.api_key = api_key or self._get_api_key_for_provider(self.provider)

        if not self.api_key:
            raise ValueError(
                f"API key not set for provider '{self.provider}'. "
                "Set the corresponding environment variable."
            )

        self.client = self._create_client(self.provider, self.api_key)
        logger.debug(f"AI client initialized for provider: {self.provider} using model: {self.model}")

    def _get_api_key_for_provider(self, provider):
        provider = provider.lower()
        if provider == "grok":
            return GROQ_API_KEY
        if provider == "google":
            return GOOGLE_API_KEY
        if provider == "deepseek":
            return DEEPSEEK_API_KEY
        if provider == "openai":
            return OPENAI_API_KEY
        if provider == "huggingface":
            return HUGGINGFACE_API_KEY
        return None

    def _create_client(self, provider, api_key):
        """Create and return the appropriate AI client."""
        provider = provider.lower()
        if provider == "grok":
            return Groq(api_key=api_key)
        if provider == "openai":
            if OpenAI is None:
                raise ImportError(
                    "openai package is not installed. Install it with `pip install openai`."
                )
            return OpenAI(api_key=api_key)
        if provider == "google":
            return GoogleAIClient(api_key=api_key)
        if provider == "huggingface":
            return InferenceClient(provider=HF_INFERENCE_PROVIDER, api_key=api_key, timeout=180)
        if provider == "deepseek":
            raise NotImplementedError(
                "Deepseek provider support is not implemented yet. "
                "Please add a Deepseek client integration."
            )
        raise ValueError(f"Unsupported AI provider: {provider}")

    def _chat_complete(self, model, messages, temperature, max_tokens):
        """
        Wraps chat completions. Uses streaming for HuggingFace to avoid
        CloudFront 504s on long generations — assembles chunks into a
        single response object identical to the non-streaming shape.
        """
        if self.provider == "huggingface":
            stream = self.client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=True,
            )
            content = ""
            finish_reason = "stop"
            for chunk in stream:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if delta.content:
                    content += delta.content
                if chunk.choices[0].finish_reason:
                    finish_reason = chunk.choices[0].finish_reason
            from types import SimpleNamespace
            return SimpleNamespace(
                choices=[SimpleNamespace(
                    message=SimpleNamespace(content=content),
                    finish_reason=finish_reason,
                )]
            )
        return self.client.chat.completions.create(
            model=model,
            messages=messages,
            response_format={"type": "json_object"},
            temperature=temperature,
            max_tokens=max_tokens,
        )

    def count_tokens(self, text):
        """Count tokens in text using tiktoken."""
        try:
            encoding = tiktoken.encoding_for_model("gpt-4o")
            tokens = encoding.encode(text)
            return len(tokens)
        except Exception as e:
            logger.warning(f"Failed to count tokens: {e}")
            return None

    def rewrite_content(self, title, body, original_text_url, seo_focus_words, sources=None, model=None):
        """
        Rewrite article content using configured AI provider.
        
        Args:
            title (str): Original article title
            body (str): Original article body in HTML
            original_text_url (str): Original URL of the article
            seo_focus_words (list): List of SEO focus words
            model (str, optional): Model name to override the default model
            
        Returns:
            dict: Contains 'title' and 'body' keys
            
        Raises:
            Exception: If API call fails or response is invalid
        """
        model_to_use = model if model else self.model
        try:
            logger.info(f"Starting content rewrite with {self.provider.upper()} using model: {model_to_use}")
            
            # Count tokens in original body
            token_count = self.count_tokens(body)
            if token_count:
                logger.debug(f"Original content: {token_count} tokens")

            prompt = self._build_rewrite_prompt(title, body, original_text_url, seo_focus_words, sources=sources)
            logger.debug(f"Sending rewrite request to {self.provider.upper()} API")

            response = self._chat_complete(
                model=model_to_use,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                max_tokens=15000,
            )


            response_text = response.choices[0].message.content

            logger.debug(f"Received response from {self.provider.upper()} API")

            # Parse JSON response (strip markdown fences if model wrapped output)
            cleaned = response_text.strip()
            if cleaned.startswith("```"):
                lines = cleaned.splitlines()
                cleaned = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
            try:
                result = json.loads(cleaned)
            except json.JSONDecodeError as e:
                logger.error(f"Failed to parse {self.provider.upper()} response as JSON: {e}")
                logger.debug(f"Response text: {response_text[:500]}")
                raise

            # Validate response structure
            required_keys = {"title", "body"}
            if not all(key in result for key in required_keys):
                missing = required_keys - set(result.keys())
                raise ValueError(f"Missing required keys in response: {missing}")

            new_title = result["title"]
            new_body = result["body"]

            logger.info(f"Content rewritten successfully")
            logger.debug(f"New title: {new_title}")
            logger.debug(f"New body length: {len(new_body)} chars")

            return {
                "title": new_title,
                "body": new_body,
            }

        except Exception as e:
            logger.error(f"Failed to rewrite content: {e}")
            raise

    def generate_seo_data(self, title, body, model=None):
        """
        Generate SEO metadata using configured AI provider.

        Args:
            title (str): Article title (rewritten)
            body (str): Article body in HTML (rewritten)
            model (str, optional): Model name to override the default model

        Returns:
            dict: Contains 'focus_keyword', 'secondary_keywords', 'meta_description',
                  'slug', 'tags', 'categories' keys
        """
        model_to_use = model if model else self.model
        try:
            logger.info(f"Generating SEO data with {self.provider.upper()} using model: {model_to_use}")

            prompt = self._build_seo_prompt(title, body)

            response = self._chat_complete(
                model=model_to_use,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                max_tokens=8000,
            )

            response_text = response.choices[0].message.content
            finish_reason = getattr(response.choices[0], "finish_reason", "unknown")
            logger.info(f"SEO raw response: finish_reason={finish_reason}, length={len(response_text) if response_text else 0}")
            logger.info(f"SEO raw response text: {repr(response_text[:800])}")

            if not response_text:
                raise ValueError(f"Empty response from {self.provider.upper()} for SEO generation")

            # Strip markdown code fences if model wrapped JSON in ```json ... ```
            cleaned = response_text.strip()
            if cleaned.startswith("```"):
                lines = cleaned.splitlines()
                cleaned = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
                logger.info(f"Stripped markdown fences; cleaned length={len(cleaned)}")

            try:
                result = json.loads(cleaned)
            except json.JSONDecodeError as e:
                # DeepSeek reasoning models emit chain-of-thought before JSON; try to extract it
                import re
                json_match = re.search(r'\{.*\}', cleaned, re.DOTALL)
                if json_match:
                    try:
                        result = json.loads(json_match.group(0))
                        logger.info("Extracted JSON from reasoning-prefixed response")
                    except json.JSONDecodeError:
                        logger.error(f"Failed to parse SEO response as JSON: {e}")
                        logger.error(f"Full SEO response (repr): {repr(response_text)}")
                        raise e
                else:
                    logger.error(f"Failed to parse SEO response as JSON: {e}")
                    logger.error(f"Full SEO response (repr): {repr(response_text)}")
                    raise

            required_keys = {"focus_keyword", "secondary_keywords", "meta_description", "slug", "tags", "categories"}
            if not all(key in result for key in required_keys):
                missing = required_keys - set(result.keys())
                raise ValueError(f"Missing required SEO keys in response: {missing}")

            logger.info(f"SEO data generated: focus_keyword='{result['focus_keyword']}'")
            logger.debug(f"SEO slug: {result['slug']}, categories: {result['categories']}")
            return result

        except Exception as e:
            logger.error(f"Failed to generate SEO data: {e}")
            raise

    @staticmethod
    def _build_seo_prompt(title, body):
        """Build the prompt for SEO data generation."""
        content = f"Title: {title}\n\n{body}"
        template = _load_prompt_template("seo_prompt.txt")
        return template.format(content=content[:MAX_CONTENT_PREVIEW])
    # @staticmethod
    # def _build_rewrite_prompt(title, body):
    #     """Build the prompt for content rewriting."""
    #     return f"""You are a professional health and supplement blog editor. Your job is to REWRITE and RESTRUCTURE content to match professional blog standards and not just summarize.

    # STRICT RULES FOR CONTENT:
    # - Do NOT shorten the article - keep all key information
    # - Do NOT skip sections, points or paragraphs
    # - Use a natural, conversational, human tone
    # - The rewritten body must be at least as long as the original

    # ### CRITICAL CITATION RULES:
    # The original content contains numbered citations (e.g., [1], [2]) and may or may not provide direct URLs. 
    # 1. **Extraction:** Identify every research study, clinical trial, or data point mentioned in the original text.
    # 2. **Format:** If URL is present, then inline citation must be a Markdown hyperlink: [1](URL) and the corresponding entry in the "References" section at the end must include the source name and hyperlink.
    #     If URL is not present then inline citation can just be a reference like [1] with no hyperlink. Corresponding entry in the "References" section  must include the source name as is from original text. Do not add any new citations that are not present in the original text.
    #     Do not remove any citations that are present in the original text. 
    # 3. **Consistency:** Ensure the numbered citations in the body match the "References" section at the end.
    # 4. **Limits:** Include up to 10 sources in the References section, but keep the inline citations focused on the 5-6 most impactful points.


    # REQUIRED HTML FORMATTING:
    # - Use proper heading hierarchy: <h2> for main sections, <h3> for subsections
    # - Wrap paragraphs in <p> tags
    # - Use <ul><li> for bullet points
    # - Use <strong> for emphasis on key terms
    # - Use <table> with <tr>, <th>, <td> for data when appropriate
    # - Image size smaller
    # - Manrope as font
    # - For hyperlinking citations, [1] should be formatted as: <a href="URL">[1]</a>


    # STRICT RULES FOR IMAGES
    # - Image not from any brand.
    # - Avif or webp images instead of jpg or png

    # STRUCTURE YOUR OUTPUT:
    # 1. Opening paragraph with engaging hook
    # 2. H2 section explaining the topic
    # 3. H2 "Why is this important?" or similar context section  
    # 4. Multiple H3 subsections for different benefits/aspects (use numbered subsections like "1. Benefit Name", "2. Next Benefit")
    # 5. H2 "Safety & Side Effects" section
    # 6. H2 "Dosage" section (if relevant)
    # 7. H2 "Bottom Line" or "Conclusion"
    # 8. H2 "Follow us for more health and supplement tips." with social links block
    # 9. H2 "References" section with all sources extracted from the original article


    # MANDATORY SOCIAL LINKS BLOCK (paste verbatim at the end of every post):
    # <!-- wp:social-links {{"openInNewTab":true,"className":"is-style-logos-only"}} -->
    # <ul class="wp-block-social-links is-style-logos-only items-justified-center" style="font-size: 4rem; gap: 1.5rem;">
    # <!-- wp:social-link {{"url":"https://x.com/SuppFactsDaily","service":"x"}} /-->
    # <!-- wp:social-link {{"url":"https://www.facebook.com/people/Suppfactsdaily/61576444959569/","service":"facebook"}} /-->
    # <!-- wp:social-link {{"url":"https://www.instagram.com/suppfactsdaily/","service":"instagram"}} /-->
    # <!-- wp:social-link {{"url":"https://www.youtube.com/@SuppFactsDaily","service":"youtube"}} /-->
    # </ul>
    # <!-- /wp:social-links -->

    # SPECIAL INSTRUCTIONS:
    # - Do NOT include author names, bios, or credits from original website
    # - For inline citations within content, format as: A study showed that [1]
    # - For the References section, use HTML anchor tags:
    # Format: [1] Source Name: <a href="link_to_article">Title or Description of article</a>
    # Example: [5] PubMed Central: <a href="https://pubmed.ncbi.nlm.nih.gov/12345678">Title of research paper</a>
    # - Make it scannable with good use of headings and bullets
    # - Include practical, actionable information

    # Return ONLY valid JSON with no extra text:
    # {{
    # "title": "A new engaging title in max 10 words",
    # "body": "Full rewritten article in proper HTML with h2, h3, p, ul, li, strong tags and numbered citations [1], [2] etc.",
    # "tags": ["tag1", "tag2", "tag3"]
    # }}

    # Original title: {title}

    # Original body:
    # {body[:MAX_CONTENT_PREVIEW]}
    # """




    @staticmethod
    def _build_rewrite_prompt(title, body, original_text_url, seo_focus_words, sources=None):
        """Build the prompt for content rewriting."""
        seo = ', '.join(seo_focus_words) if seo_focus_words else 'None'
        if sources:
            sources_html = "\n<h2>Sources</h2>\n<ul>\n"
            for s in sources:
                url = s.get("url", "")
                text = s.get("text", "")
                if url:
                    sources_html += f'  <li><a href="{url}">{text}</a></li>\n'
                else:
                    sources_html += f'  <li>{text}</li>\n'
            sources_html += "</ul>"
            body = body + sources_html
        template = _load_prompt_template("rewrite_prompt.txt")
        return template.format(
            title=title,
            body=body[:MAX_CONTENT_PREVIEW],
            original_text_url=original_text_url,
            seo_focus_words=seo,
        )