import sys
import json
import argparse
import logging
import os
import re
import time
from datetime import datetime


class ColoredFormatter(logging.Formatter):
    # ANSI escape codes for colors
    COLORS = {
        "DEBUG": "\033[94m",    # Blue
        "INFO": "\033[92m",     # Green
        "WARNING": "\033[93m",  # Yellow
        "ERROR": "\033[91m",    # Red
        "CRITICAL": "\033[95m", # Magenta
        "RESET": "\033[0m"      # Reset
    }

    def format(self, record):
        color = self.COLORS.get(record.levelname, self.COLORS["RESET"])
        message = super().format(record)
        return f"{color}{message}{self.COLORS['RESET']}"

def setup_logger(log_file=None, level=logging.DEBUG):
    """
    Sets up a logger that logs messages to the console and optionally to a file.

    Parameters:
        log_file (str): Path to a log file. If None, only console logging is used.
        level (int): Logging level (e.g., logging.INFO, logging.DEBUG, logging.ERROR).

    Returns:
        logging.Logger: Configured logger instance.
    """
    logger = logging.getLogger(__name__)
    logger.setLevel(level)

    # Avoid adding multiple handlers if logger already exists
    if not logger.handlers:
        # Console handler (INFO and above)
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)
        ch_formatter = ColoredFormatter('%(asctime)s - %(levelname)s - %(message)s')
        ch.setFormatter(ch_formatter)
        logger.addHandler(ch)

        # File handler (optional)
        if log_file:
            # Ensure the directory exists
            os.makedirs(os.path.dirname(log_file), exist_ok=True)
            fh = logging.FileHandler(log_file, mode='a')
            fh.setLevel(level)
            fh_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
            fh.setFormatter(fh_formatter)
            logger.addHandler(fh)

    return logger

# -------- OpenAI --------
from openai import OpenAI

# -------- Gemini --------
from google import genai

# -------- Hugging Face --------
from huggingface_hub import InferenceClient
import requests

# -------------------------
# API Keys
# -------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") or "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"
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or "AIzaSyDRO5QyNa44eUsbfE5eQbMpjagO3398Ihs"

# -------------------------
# Hugging Face
# -------------------------
HF_API_KEY = os.getenv("HF_API_KEY") or "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"
HF_MODEL = "Qwen/Qwen2.5-7B-Instruct"
HF_API_URL = f"https://api-inference.huggingface.co/models/{HF_MODEL}"

# -------------------------
# Groq
# -------------------------
# export GROQ_API_KEY= "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"
GROQ_API_KEY = os.getenv("GROQ_API_KEY") or "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"
GROQ_MODEL = "llama-3.1-8b-instant"
GROQ_MODEL = "openai/gpt-oss-120b"
# -------------------------
# Configuration
# -------------------------
OPENAI_MODEL = "gpt-4.1-mini"
GEMINI_MODEL = "gemini-2.0-flash"
REQUIRED_NUTRITION_FIELDS = [
"calories",
"proteinContent",
"fatContent",
"carbohydrateContent"]

REQUIRED_FIELDS = [
    "keywords",
    "dateCreated",
    "recipeYield",
    "recipeIngredient",
    "recipeInstructions"
]


# -------------------------
# Prompt generation
# -------------------------
def generate_prompt(recipe_content: str) -> str:
    return f"""
Du bist eine reine JSON-Transformationsmaschine.

Anweisungen:

1. Aktualisiere oder erstelle ein Objekt „nutrition“ für das eine Portion des Rezepts im JSON und berechne die folgenden Felder:
   "nutrition": {{
        "@type": "NutritionInformation",
        "servingSize": "1 Portion",
        "calories": "<numeric> kcal",
        "proteinContent": "<numeric> g",
        "fatContent": "<numeric> g",
        "carbohydrateContent": "<numeric> g"                                      
        "cholesterolContent": "<numeric> g",
        "fiberContent": "<numeric> g",
        "saturatedFatContent": "<numeric> g",
        "sodiumContent": "<numeric> g",
        "sugarContent": "<numeric> g",
        "transFatContent": "<numeric> g",
        "unsaturatedFatContent": "<numeric> g"
   }}
    Berechne die Nährwerte anhand der Zutaten in „recipeIngredient“.
    Die Zutaten sind auf Deutsch.
    Füge die Einheiten zu den Feldern hinzu.
    Falls du Annahmen bezüglich Mengen machst, schreibe diese zu den Zutaten, recipeIngredient.
2. Alle anderen Felder im ursprünglichen JSON unverändert lassen.
3. Gib das vollständige aktualisierte JSON-Objekt zurück, einschließlich der aktualisierten Nährwerte.
4. Gib ausschließlich gültiges JSON zurück.
5. Keine Markdown-Formatierung, keine Kommentare, keine Erklärungen und keinen zusätzlichen Text ausgeben.
6. Die Ausgabe muss direkt mit json.loads() parsbar sein.

Data:
{recipe_content}
"""
# -------------------------
# Hugging Face call (LEGACY – WORKING)
# -------------------------
def call_huggingface(prompt: str) -> str:
    API_URL = "https://router.huggingface.co/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {HF_API_KEY}",
    }

    def query(payload):
        response = requests.post(API_URL, headers=headers, json=payload)
        return response.json()

    response = query({
        "messages": [
            {
                "role": "user",
                "content": prompt
            }
        ],
        "model": "Qwen/Qwen2.5-72B:featherless-ai"
    })

    returnStr = response["choices"][0]["message"]["content"]
    logger.debug(response)
    logger.debug(returnStr)
    return returnStr

# -------------------------
# Groq call
# -------------------------
def call_groq(prompt: str) -> str:
    client = OpenAI(
        api_key=GROQ_API_KEY,
        base_url="https://api.groq.com/openai/v1"
    )

    response = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {"role": "system", "content": "You return valid JSON only."},
            {"role": "user", "content": prompt},
        ],
        temperature=0.3,
    )
    logger.debug(response)
    returnStr = response.choices[0].message.content.strip()
    return returnStr

# -------------------------
# OpenAI call
# -------------------------
def call_openai(prompt: str) -> str:
    client = OpenAI(api_key=OPENAI_API_KEY)

    response = client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": "You return valid JSON only."},
            {"role": "user", "content": prompt},
        ],
        temperature=0.3,
    )

    return response.choices[0].message.content.strip()


# -------------------------
# Gemini call
# -------------------------
def call_gemini(prompt: str) -> str:
    client = genai.Client(api_key=GEMINI_API_KEY)

    response = client.models.generate_content(
        model=GEMINI_MODEL,
        contents=prompt,
    )

    text = response.text.strip()

    # Defensive cleanup
    if text.startswith("```"):
        text = text.split("```")[1].replace("json", "").strip()

    return text
# -------------------------
# Model router
# -------------------------
def call_model(provider: str, prompt: str) -> str:
    if provider == "openai":
        return call_openai(prompt)
    elif provider == "gemini":
        return call_gemini(prompt)
    elif provider == "huggingface":
        return call_huggingface(prompt)
    elif provider == "groq":
        return call_groq(prompt)
    else:
        raise ValueError(f"Unknown provider: {provider}")
# -------------------------
# Evaluation
# -------------------------

def list_dict_fields(d, parent_key=""):
    """
    Recursively lists all keys and subkeys in a dict, returning full paths.
    """
    fields = []

    for k, v in d.items():
        full_key = f"{parent_key}.{k}" if parent_key else k

        fields.append(full_key)

        if isinstance(v, dict):
            fields.extend(list_dict_fields(v, full_key))
        elif isinstance(v, list):
            # If list of dicts, include their fields
            for i, item in enumerate(v):
                if isinstance(item, dict):
                    fields.extend(list_dict_fields(item, f"{full_key}[{i}]"))

    return fields

def evaluate_response(data: dict) -> dict:
    errors = []
    all_fields = list_dict_fields(data)
    for field in  REQUIRED_FIELDS + REQUIRED_NUTRITION_FIELDS :
        if not any(field in fields for fields in all_fields):
            errors.append(f"Missing field: {field}")

    # Check that recipeYield is numeric
    recipe_yield = data.get("recipeYield")
    if not isinstance(recipe_yield, (int, float)):
        errors.append("recipeYield must be numeric")

    return {
        "valid": len(errors) == 0,
        "errors": errors
    }


# -------------------------
# Main processing for a single file
# -------------------------

def remove_keys(d: dict, keys_to_remove: list):
    """
    Remove specified keys from a dictionary (top-level only).
    """
    for key in keys_to_remove:
        if key in d:
            del d[key]
    return d


def extract_and_remove_keys(d: dict, keys_to_remove: list):
    """
    Remove specified keys from a dictionary (top-level only) and store their values.

    Returns a tuple: (cleaned_dict, extracted_values)
    """
    extracted = {}
    for key in keys_to_remove:
        if key in d:
            extracted[key] = d.pop(key)
    return d, extracted


def process_recipe(file_path: str, provider: str, update_nutrition: bool = True, delete_nutrition: bool = True):
    if not os.path.exists(file_path):
        logger.error(f"❌ File not found: {file_path}")
        return

    with open(file_path, "r", encoding="utf-8") as f:
        try:
            recipe_data = json.load(f)
        except json.JSONDecodeError:
            logger.error(f"❌ File is not valid JSON: {file_path}")
            return

    # -------------------------
    # Nutrition check (optional)
    # -------------------------
    if "nutrition" in recipe_data and isinstance(recipe_data["nutrition"], dict):
        all_fields = list_dict_fields(recipe_data["nutrition"])
        for field in REQUIRED_NUTRITION_FIELDS:
            if not any(field in fields for fields in all_fields):
                update_nutrition = True
                logger.warning(f"Missing field: {field}, Nutrition block will be updated.")
                break
        warning = f"⚠️ Nutrition block already exists in {os.path.basename(os.path.dirname(file_path))}/{os.path.basename(file_path)}. "
        if update_nutrition:
            logger.warning(f"{warning:<100} - It will be updated")
        elif delete_nutrition:
            logger.warning(f"{warning:<100} - It will be deleted")
            recipe_data, _ = extract_and_remove_keys(recipe_data, ["nutrition"])
            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(recipe_data, f, indent=2, ensure_ascii=False)
            return True
        else:
            logger.warning(f"{warning:<100} - Skipping recipe")
            return True

    else:
        logger.info(f"ℹ️ No nutrition block found in {file_path}. A new one will be added.")

    recipe_data,_ = extract_and_remove_keys(recipe_data, ["video","aggregateRating","reviews"])
    recipe_data, stored_keys = extract_and_remove_keys(recipe_data, ["valueInit"])


    # Serialize JSON back to string for prompt
    recipe_content = json.dumps(recipe_data, ensure_ascii=False)
    prompt = generate_prompt(recipe_content)

    logger.info(f"🚀 Using provider: {provider} → {file_path}")

    logger.debug(prompt)

    while(True):
        try:
            raw_response = call_model(provider, prompt)
            start = raw_response.find("{")
            end = raw_response.rfind("}")
            raw_response = raw_response[start:end + 1]
            updated_data = json.loads(raw_response)

            evaluation = evaluate_response(updated_data)
            if not evaluation["valid"]:
                logger.error("❌ Evaluation failed. File NOT updated.")
                for e in evaluation["errors"]:
                    logger.error(f"{e}")
                return

            combined = updated_data | stored_keys

            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(combined, f, indent=2, ensure_ascii=False)

            logger.info("✅ Recipe updated successfully.")
            break

        except json.JSONDecodeError:
            logger.error("❌ Model did not return valid JSON.")
            logger.error("RAW OUTPUT:\n", raw_response)

        except Exception as e:
            msg = str(e)
            if "rate limit" in msg.lower():
                # Extract wait time
                match = re.search(r'Please try again in (\d+)m([\d\.]+)s', msg)
                logger.warning(msg)
                if match:
                    minutes = int(match.group(1))
                    seconds = float(match.group(2))
                    wait_time = (minutes * 60 + seconds) *1.5
                    logger.warning(f"Rate limit hit. Waiting for {wait_time:.2f} seconds...")
                    for remaining in range(int(wait_time), 0, -1):
                        print(f"\rTime left: {remaining}s", end="")
                        time.sleep(1)
                    logger.info("\nRetrying now...")
                else:
                    logger.error("Rate limit hit, could not parse wait time. Waiting 200 seconds...")
                    return False

            else:
                logger.error("❌ Error:", e)
                return False
    return True

# -------------------------
# Recursively find all recipe.json files in a folder
# -------------------------
def process_folder(folder_path: str, provider: str, update_nutrition: bool = True , delete_nutrition: bool = True):
    json_files = []

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower() == "recipe.json":
                full_path = os.path.join(root, file)
                json_files.append(os.path.join(root, file))
    i = 1
    for file in json_files:
        logger.info(f"Checking: {i}/{len(json_files)} ... {os.path.basename(os.path.dirname(file))}/{os.path.basename(file)}.")
        i+=1
        if not process_recipe(file, provider, update_nutrition, delete_nutrition):
            logger.warning(f"File: {file} failed. Continuing")

# -------------------------
# CLI entry point
# -------------------------
if __name__ == "__main__":
    script_name = os.path.splitext(os.path.basename(__file__))[0]

    log_path = os.path.join(os.path.expanduser("~"), ".LOG",f"{script_name}",
                            f"{script_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
                            )
    logger = setup_logger(log_file=log_path)
    logger.info("Logger initialized.")

    parser = argparse.ArgumentParser()
    parser.add_argument("path", help="Recipe JSON file or folder containing recipe.json files")
    parser.add_argument(
        "--provider",
        choices=["openai", "gemini", "huggingface", "groq"],
        default="groq",
        help="LLM provider to use"
    )
    parser.add_argument(
        "-un", "--update_nutrition",
        action="store_true",
        help="Enable nutrition block check and update"
    )
    parser.add_argument(
        "-dn", "--delete_nutrition",
        action="store_true",
        help="Enable nutrition block deletion"
    )

    args = parser.parse_args()
    path = args.path
    provider = args.provider
    update_nutrition = args.update_nutrition
    delete_nutrition = args.delete_nutrition
  #  path = "/data/Nextcloud/Documents/Dokumente/Shared/Diverses/Kochbuch/RecepiesLucas/Sauerteig herstellen und züchten"
#    update_nutrition = False
#    delete_nutrition = False

    if os.path.isdir(path):
        process_folder(path, provider, update_nutrition, delete_nutrition)
    else:
        process_recipe(path, provider, update_nutrition, delete_nutrition)