import os
import re
import zipfile
import tempfile
from pathlib import Path
from uuid import uuid4
from typing import List, Tuple

import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv(override=True)

# ---------- CLIENT SETUP ----------

groq_api_key = os.getenv('GROQ_API_KEY')
if groq_api_key:
    print(f"Groq API Key exists and begins {groq_api_key[:4]}")
else:
    print("Groq API Key not set")

client = OpenAI(api_key=groq_api_key, base_url="https://api.groq.com/openai/v1")
MODEL = "openai/gpt-oss-20b"

# ---------- PROMPTS ----------

SYSTEM_PROMPT = """
You are a synthetic data generator.

You will be given:
- A schema description
- A target format (e.g., json, jsonl, csv)
- A number of records to generate in this batch

You MUST generate multiple files in a strict, machine-parsable structure.

For EACH file, output:

===FILE_START===
FILENAME: <safe_file_name>
<file content>
===FILE_END===

Rules:
- FILENAME must be derived from the PRIMARY IDENTIFIER field in the record itself.
  For example: if the record has a "name" field with value "Sarah Johnson", the filename must be "sarah_johnson.json".
  If the record has a "title" field with value "Service Contract Alpha", use "service_contract_alpha.json".
  If the record has a "product_name" field with value "Blue Widget Pro", use "blue_widget_pro.json".
- NEVER copy field values from the schema description or examples verbatim.
- All values must be freshly generated and diverse.
- Names, roles, companies must vary significantly across files.
- NEVER use generic names like employee_001, record_01, file_1, or any numbered fallback.
- FILENAME must use only lowercase letters, numbers, and underscores (replace spaces with underscores).
- FILENAME must include the correct extension (e.g., .json, .jsonl, .csv).
- No spaces or special characters in FILENAME.
- File content must follow the schema and format.
- No explanations, no markdown, no comments.
- Do NOT output anything outside the specified structure.
"""


def build_batch_prompt(schema: str, file_type: str, records_in_batch: int) -> str:
    return f"""
Generate {records_in_batch} synthetic files.

Schema:
{schema}

Target format: {file_type}

For EACH file, output:

===FILE_START===
FILENAME: <safe_file_name>
<file content>
===FILE_END===

Rules:
- FILENAME must come directly from the primary identifier inside the record.
  Example: if the record contains {{"full_name": "Marcus Rivera"}}, the filename is marcus_rivera.{file_type}
  Example: if the record contains {{"contract_title": "NDA Agreement Beta"}}, the filename is nda_agreement_beta.{file_type}
- NEVER use numbered placeholders like employee_001 or record_02.
- Each filename must be unique and reflect its specific content.
- No explanations, no extra text outside the structure.
"""


def call_llm_for_batch(schema: str, file_type: str, records_in_batch: int) -> str:
    prompt = build_batch_prompt(schema, file_type, records_in_batch)
    try:
        resp = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.6,
        )
        return resp.choices[0].message.content
    except Exception as e:
        raise gr.Error(f"LLM call failed: {e}")


def parse_batch_output(text: str) -> List[Tuple[str, str]]:
    results: List[Tuple[str, str]] = []
    blocks = text.split("===FILE_START===")
    for block in blocks[1:]:
        if "===FILE_END===" not in block:
            continue
        file_part, _ = block.split("===FILE_END===", 1)
        lines = [ln for ln in file_part.strip().splitlines() if ln.strip()]
        if not lines:
            continue
        first_line = lines[0].strip()
        if first_line.startswith("FILENAME:"):
            filename = first_line.replace("FILENAME:", "").strip()
            content = "\n".join(lines[1:]).strip()
        else:
            filename = f"unnamed_{uuid4().hex}.txt"
            content = "\n".join(lines).strip()
        results.append((filename, content))
    return results


def ensure_unique_filename(filename: str, used: set) -> str:
    base = filename
    if base in used:
        stem, dot, ext = base.rpartition(".")
        stem = stem or "file"
        ext = ("." + ext) if dot else ""
        i = 1
        new_name = f"{stem}_{i}{ext}"
        while new_name in used:
            i += 1
            new_name = f"{stem}_{i}{ext}"
        filename = new_name
    used.add(filename)
    return filename


def generate_files(schema: str, file_type: str, num_files: int, batch_size: int, progress=gr.Progress()) -> Tuple[str, str]:
    if not schema.strip():
        raise gr.Error("Please provide a schema / structure description.")
    if num_files < 1:
        raise gr.Error("Number of files must be at least 1.")
    if batch_size < 1:
        raise gr.Error("Batch size must be at least 1.")

    total_batches = (num_files + batch_size - 1) // batch_size
    tmp_dir = Path(tempfile.mkdtemp())
    zip_path = tmp_dir / "synthetic_data.zip"
    used_filenames = set()
    files_written = 0
    preview_files = []  # store first 3 for preview

    progress(0, desc="Starting generation...")

    with zipfile.ZipFile(zip_path, "w") as z:
        for batch_index in range(total_batches):
            remaining = num_files - files_written
            current_batch_size = min(batch_size, remaining)

            progress(
                batch_index / total_batches,
                desc=f"Batch {batch_index + 1}/{total_batches} — generating {current_batch_size} files..."
            )

            batch_output = call_llm_for_batch(schema, file_type, current_batch_size)
            parsed = parse_batch_output(batch_output)

            for filename, content in parsed:
                if files_written >= num_files:
                    break
                filename = ensure_unique_filename(filename, used_filenames)
                file_path = tmp_dir / filename
                file_path.write_text(content, encoding="utf-8")
                z.write(file_path, arcname=filename)

                # Collect first 3 files for preview
                if len(preview_files) < 3:
                    preview_files.append((filename, content))

                files_written += 1

            if files_written >= num_files:
                break

    progress(1.0, desc=f"Done! {files_written} files generated.")

    # Build preview text
    preview_text = ""
    for fname, fcontent in preview_files:
        preview_text += f"### 📄 {fname}\n```\n{fcontent}\n```\n\n"

    if files_written > 3:
        preview_text += f"*...and {files_written - 3} more files in the ZIP.*"

    return str(zip_path), preview_text


def ui_generate(schema, file_type, num_files, batch_size):
    return generate_files(schema, file_type, num_files, batch_size)


# ---------- GRADIO UI ----------

with gr.Blocks(title="Synthetic Data Generator") as demo:
    gr.Markdown("# Synthetic Data Generator")
    gr.Markdown(
        "Generate synthetic files using your schema. "
        "Filenames are created by the LLM and uniqueness is enforced locally."
    )

    schema_input = gr.Textbox(
        label="Schema / structure description",
        lines=10,
        placeholder='Describe your data schema in plain English. Example:\n"Generate synthetic employee records with fields: full_name, role, department, salary, hire_date. Use diverse, realistic values.',
    )
    file_type_input = gr.Dropdown(
        choices=["json", "jsonl", "csv", "txt", "md"],
        value="json",
        label="File type",
    )
    num_files_input = gr.Slider(minimum=1, maximum=200, value=20, step=1, label="Number of files to generate")
    batch_size_input = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Files per LLM call (batch size)")

    generate_btn = gr.Button("Generate", variant="primary")

    with gr.Row():
        zip_output = gr.File(label="Download ZIP")

    preview_output = gr.Markdown(label="Preview (first 3 files)")

    generate_btn.click(
        fn=ui_generate,
        inputs=[schema_input, file_type_input, num_files_input, batch_size_input],
        outputs=[zip_output, preview_output],
    )

if __name__ == "__main__":
    demo.launch(inbrowser=True)