import os
import re
import json
import duckdb
import faiss
import numpy as np

from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel

from sentence_transformers import SentenceTransformer
from groq import Groq


# ============================================================
# 1. DATABASE
# ============================================================

con = duckdb.connect(":memory:")

con.execute("""
CREATE TABLE sales (
    sale_id INTEGER,
    region VARCHAR,
    store VARCHAR,
    category VARCHAR,
    product VARCHAR,
    quantity INTEGER,
    revenue DOUBLE,
    sale_date DATE
)
""")

con.execute("""
INSERT INTO sales VALUES
    (1,'Riyadh','Store A','Dairy','Milk 1L',120,600.0,'2026-07-05'),
    (2,'Riyadh','Store A','Bakery','Bread',80,240.0,'2026-07-05'),
    (3,'Jeddah','Store B','Dairy','Milk 1L',95,475.0,'2026-07-06'),
    (4,'Jeddah','Store B','Produce','Tomatoes',200,300.0,'2026-07-06'),
    (5,'Dammam','Store C','Dairy','Milk 1L',60,300.0,'2026-07-07'),
    (6,'Riyadh','Store A','Produce','Tomatoes',150,225.0,'2026-08-01'),
    (7,'Jeddah','Store B','Bakery','Bread',110,330.0,'2026-08-01')
""")


SCHEMA_DESC = """
Table: sales

Columns:
sale_id INTEGER
region VARCHAR
store VARCHAR
category VARCHAR
product VARCHAR
quantity INTEGER
revenue DOUBLE
sale_date DATE
"""


# ============================================================
# 2. TRACE CORPUS
# ============================================================

TRACE_LIBRARY = [

    {
        "question": "What is total revenue by region for July 2026?",
        "trace": {
            "intent": "aggregation",
            "metric": "revenue",
            "aggregation": "SUM",
            "dimension": "region",
            "time_filter": "July 2026",
            "grouping": "region",
            "ordering": None,
            "limit": None
        },
        "sql": """
SELECT
    region,
    SUM(revenue) AS total_revenue
FROM sales
WHERE sale_date >= '2026-07-01'
  AND sale_date < '2026-08-01'
GROUP BY region;
"""
    },

    {
        "question": "Show total quantity sold by category for August 2026.",
        "trace": {
            "intent": "aggregation",
            "metric": "quantity",
            "aggregation": "SUM",
            "dimension": "category",
            "time_filter": "August 2026",
            "grouping": "category",
            "ordering": None,
            "limit": None
        },
        "sql": """
SELECT
    category,
    SUM(quantity) AS total_quantity
FROM sales
WHERE sale_date >= '2026-08-01'
  AND sale_date < '2026-09-01'
GROUP BY category;
"""
    },

    {
        "question": "Which store had the highest revenue overall?",
        "trace": {
            "intent": "ranking",
            "metric": "revenue",
            "aggregation": "SUM",
            "dimension": "store",
            "time_filter": None,
            "grouping": "store",
            "ordering": "total revenue DESC",
            "limit": 1
        },
        "sql": """
SELECT
    store,
    SUM(revenue) AS total_revenue
FROM sales
GROUP BY store
ORDER BY total_revenue DESC
LIMIT 1;
"""
    }
]


# ============================================================
# 3. EMBEDDING MODEL
# ============================================================

embedder = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2"
)


# ============================================================
# 4. FAISS TRACE INDEX
# ============================================================

def build_index():

    questions = [item["question"] for item in TRACE_LIBRARY]

    embeddings = embedder.encode(questions, normalize_embeddings=True)
    embeddings = np.asarray(embeddings, dtype="float32")

    index = faiss.IndexFlatIP(embeddings.shape[1])
    index.add(embeddings)

    return index


index = build_index()

# Minimum cosine similarity (embeddings are normalized, so FAISS IP
# score == cosine similarity) for a retrieved trace to be considered
# "relevant" enough to inject into the prompt as a verified example.
#
# 0.55 is a reasonable starting cutoff for all-MiniLM-L6-v2 on short
# question-style text: near-duplicate/paraphrased questions usually
# score 0.7+, loosely related questions land around 0.4-0.6, and
# unrelated questions fall below 0.3. Tune against your own
# TRACE_LIBRARY as it grows -- it's a heuristic, not a guarantee.
MIN_SIMILARITY_SCORE = 0.55


def retrieve_traces(question, k=3, min_score=MIN_SIMILARITY_SCORE):

    if not TRACE_LIBRARY:
        return []

    k = min(k, len(TRACE_LIBRARY))

    query_embedding = embedder.encode([question], normalize_embeddings=True)
    query_embedding = np.asarray(query_embedding, dtype="float32")

    scores, indices = index.search(query_embedding, k)

    results = []

    for score, idx in zip(scores[0], indices[0]):

        if idx < 0:
            continue

        score = float(score)

        # Below-threshold matches are dropped rather than passed to
        # the LLM as "verified" -- a weak match is worse than no
        # match, since it can anchor the model to the wrong pattern.
        if score < min_score:
            continue

        item = TRACE_LIBRARY[idx].copy()
        item["score"] = score
        results.append(item)

    return results


# ============================================================
# 5. GROQ CLIENT
# ============================================================

def get_client():

    api_key = os.environ.get("GROQ_API_KEY")

    if not api_key:
        raise RuntimeError("GROQ_API_KEY is not set.")

    return Groq(api_key=api_key)


# ============================================================
# 6. FORMAT RETRIEVED TRACES
# ============================================================

def format_retrieved_traces(examples):

    output = ""

    for i, example in enumerate(examples, start=1):

        output += f"""

==============================
VERIFIED EXAMPLE {i}
Similarity: {example["score"]:.3f}
==============================

Question:
{example["question"]}

Structured Solution Trace:
{json.dumps(example["trace"], indent=2)}

Verified SQL:
{example["sql"]}

"""

    return output


# ============================================================
# 7. GENERATE STRUCTURED TRACE
# ============================================================

def generate_trace(question, retrieved_examples):

    client = get_client()
    examples = format_retrieved_traces(retrieved_examples)

    prompt = f"""
You are an expert retail Text-to-SQL planner.

Database schema:

{SCHEMA_DESC}

You have access to VERIFIED previous SQL solutions.

{examples}

New user question:

{question}

Create a structured SQL solution trace.

IMPORTANT:

Do NOT provide hidden chain-of-thought.

Only provide the following explicit fields:

intent
metric
aggregation
dimension
time_filter
grouping
ordering
limit

Use null when a field does not apply.

Return ONLY valid JSON.

Example format:

{{
  "intent": "aggregation",
  "metric": "revenue",
  "aggregation": "SUM",
  "dimension": "category",
  "time_filter": "August 2026",
  "grouping": "category",
  "ordering": null,
  "limit": null
}}
"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
        temperature=0
    )

    text = response.choices[0].message.content.strip()

    text = re.sub(r"```json\s*", "", text)
    text = re.sub(r"```\s*", "", text)

    try:
        return json.loads(text)
    except Exception:
        return {"raw_trace": text}


# ============================================================
# 8. GENERATE SQL FROM TRACE
# ============================================================

def generate_sql(question, trace, retrieved_examples):

    client = get_client()
    examples = format_retrieved_traces(retrieved_examples)

    prompt = f"""
You are an expert DuckDB Text-to-SQL generator.

Database schema:

{SCHEMA_DESC}

VERIFIED PREVIOUS EXAMPLES:

{examples}

NEW USER QUESTION:

{question}

NEW STRUCTURED SOLUTION TRACE:

{json.dumps(trace, indent=2)}

Generate SQL that correctly answers the question.

Rules:

1. Use only the provided schema.
2. Use DuckDB SQL.
3. Follow the structured trace.
4. Do not invent columns.
5. Handle date ranges correctly.
6. Return ONLY SQL.
"""

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
        temperature=0
    )

    sql = response.choices[0].message.content.strip()

    sql = re.sub(r"```sql\s*", "", sql, flags=re.IGNORECASE)
    sql = re.sub(r"```\s*", "", sql)

    return sql.strip()


# ============================================================
# 9. TRACE LIBRARY PREVIEW
# ============================================================

def library_preview():
    return [item["question"] for item in TRACE_LIBRARY]


# ============================================================
# 10. FASTAPI APP
# ============================================================

app = FastAPI(title="Text-to-SQL Trace RAG")


class AskRequest(BaseModel):
    question: str


class AddTraceRequest(BaseModel):
    question: str
    trace: dict
    sql: str


@app.post("/api/ask")
def ask(req: AskRequest):

    question = req.question.strip()

    if not question:
        return {
            "retrieved": [],
            "trace": None,
            "sql": "",
            "columns": [],
            "rows": [],
            "status": "Please enter a question."
        }

    try:
        examples = retrieve_traces(question, k=3)

        retrieved = [
            {"question": e["question"], "score": e["score"]}
            for e in examples
        ]

        trace = generate_trace(question, examples)
        sql = generate_sql(question, trace, examples)

        try:
            df = con.execute(sql).fetchdf()
            columns = list(df.columns)
            rows = df.to_dict(orient="records")
            status = "SQL executed successfully."
        except Exception as e:
            columns = []
            rows = []
            status = f"SQL execution failed: {str(e)}"

        return {
            "retrieved": retrieved,
            "trace": trace,
            "sql": sql,
            "columns": columns,
            "rows": rows,
            "status": status
        }

    except Exception as e:
        return {
            "retrieved": [],
            "trace": None,
            "sql": "",
            "columns": [],
            "rows": [],
            "status": f"Pipeline error: {str(e)}"
        }


@app.post("/api/add_trace")
def add_trace(req: AddTraceRequest):

    global index

    question = req.question.strip()
    sql = req.sql.strip()

    if not question:
        return {"status": "Question is required.", "library": library_preview()}

    if not sql:
        return {"status": "SQL is required.", "library": library_preview()}

    TRACE_LIBRARY.append({
        "question": question,
        "trace": req.trace,
        "sql": sql
    })

    # Rebuild FAISS index. Note: this mutates process-global state
    # with no locking -- fine for a single-worker prototype, not
    # safe under concurrent requests or multiple uvicorn workers.
    index = build_index()

    return {
        "status": f"Verified trace added. Trace corpus now contains {len(TRACE_LIBRARY)} examples.",
        "library": library_preview()
    }


@app.get("/api/library")
def get_library():
    return {"library": library_preview()}


# ============================================================
# 11. FRONTEND (single inline HTML page, vanilla JS, no build step)
# ============================================================

INDEX_HTML = """
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>Text-to-SQL Trace RAG</title>
<style>
  body { font-family: -apple-system, Segoe UI, Roboto, sans-serif; max-width: 900px; margin: 40px auto; padding: 0 20px; color: #1a1a1a; }
  h1 { font-size: 1.4rem; }
  .tabs { display: flex; gap: 8px; margin-bottom: 20px; }
  .tab-btn { padding: 8px 16px; border: 1px solid #ccc; background: #f5f5f5; cursor: pointer; border-radius: 6px; }
  .tab-btn.active { background: #2563eb; color: white; border-color: #2563eb; }
  .tab { display: none; }
  .tab.active { display: block; }
  textarea, input[type=text] { width: 100%; box-sizing: border-box; padding: 8px; font-family: inherit; }
  textarea { min-height: 70px; }
  button.run { margin-top: 10px; padding: 10px 18px; background: #2563eb; color: white; border: none; border-radius: 6px; cursor: pointer; }
  pre { background: #f5f5f5; padding: 12px; border-radius: 6px; overflow-x: auto; white-space: pre-wrap; }
  table { border-collapse: collapse; width: 100%; margin-top: 8px; }
  th, td { border: 1px solid #ddd; padding: 6px 10px; text-align: left; font-size: 0.9rem; }
  th { background: #f0f0f0; }
  .status { margin-top: 8px; font-size: 0.9rem; color: #444; }
  .section-title { font-weight: 600; margin-top: 18px; margin-bottom: 6px; }
</style>
</head>
<body>

<h1>Text-to-SQL Trace RAG</h1>
<p>Retrieves verified SQL solution traces (FAISS, cosine &ge; 0.55) and uses them to ground new Text-to-SQL generation.</p>

<div class="tabs">
  <button class="tab-btn active" onclick="showTab('ask', event)">Ask Question</button>
  <button class="tab-btn" onclick="showTab('add', event)">Add Verified Trace</button>
</div>

<div id="ask" class="tab active">
  <textarea id="question" placeholder="e.g. What is total revenue by category for August 2026?"></textarea>
  <button class="run" onclick="runAsk()">Generate &amp; Run</button>

  <div class="section-title">Retrieved Similar Traces</div>
  <pre id="retrieved">-</pre>

  <div class="section-title">Structured Solution Trace</div>
  <pre id="trace">-</pre>

  <div class="section-title">Generated SQL</div>
  <pre id="sql">-</pre>

  <div class="section-title">Query Result</div>
  <div id="result">-</div>

  <div class="status" id="askStatus"></div>
</div>

<div id="add" class="tab">
  <p>Only add an example after a human has verified the trace and SQL are correct and produce the expected result.</p>
  <label>Question</label>
  <input type="text" id="vQuestion">
  <label>Structured Trace (JSON)</label>
  <textarea id="vTrace">{}</textarea>
  <label>Verified SQL</label>
  <textarea id="vSql"></textarea>
  <button class="run" onclick="runAddTrace()">Add Verified Trace</button>
  <div class="status" id="addStatus"></div>

  <div class="section-title">Current Trace Corpus</div>
  <pre id="library">-</pre>
</div>

<script>
function showTab(name, evt) {
  document.querySelectorAll('.tab').forEach(el => el.classList.remove('active'));
  document.querySelectorAll('.tab-btn').forEach(el => el.classList.remove('active'));
  document.getElementById(name).classList.add('active');
  evt.target.classList.add('active');
  if (name === 'add') refreshLibrary();
}

async function runAsk() {
  const question = document.getElementById('question').value;
  document.getElementById('askStatus').textContent = 'Running...';

  const res = await fetch('/api/ask', {
    method: 'POST',
    headers: {'Content-Type': 'application/json'},
    body: JSON.stringify({question})
  });
  const data = await res.json();

  document.getElementById('retrieved').textContent =
    data.retrieved.length
      ? data.retrieved.map(r => `${r.score.toFixed(3)}  ${r.question}`).join('\\n')
      : 'No verified trace scored above the similarity threshold -- generated zero-shot (no examples injected into the prompt).';

  document.getElementById('trace').textContent = JSON.stringify(data.trace, null, 2);
  document.getElementById('sql').textContent = data.sql;

  const resultDiv = document.getElementById('result');
  if (data.columns.length) {
    let html = '<table><tr>' + data.columns.map(c => `<th>${c}</th>`).join('') + '</tr>';
    data.rows.forEach(row => {
      html += '<tr>' + data.columns.map(c => `<td>${row[c]}</td>`).join('') + '</tr>';
    });
    html += '</table>';
    resultDiv.innerHTML = html;
  } else {
    resultDiv.textContent = 'No rows returned.';
  }

  document.getElementById('askStatus').textContent = data.status;
}

async function runAddTrace() {
  const question = document.getElementById('vQuestion').value;
  const sql = document.getElementById('vSql').value;
  let trace;
  try {
    trace = JSON.parse(document.getElementById('vTrace').value);
  } catch (e) {
    document.getElementById('addStatus').textContent = 'Trace must be valid JSON.';
    return;
  }

  const res = await fetch('/api/add_trace', {
    method: 'POST',
    headers: {'Content-Type': 'application/json'},
    body: JSON.stringify({question, trace, sql})
  });
  const data = await res.json();

  document.getElementById('addStatus').textContent = data.status;
  document.getElementById('library').textContent = data.library.map((q,i) => `${i+1}. ${q}`).join('\\n');
}

async function refreshLibrary() {
  const res = await fetch('/api/library');
  const data = await res.json();
  document.getElementById('library').textContent = data.library.map((q,i) => `${i+1}. ${q}`).join('\\n');
}
</script>

</body>
</html>
"""


@app.get("/", response_class=HTMLResponse)
def home():
    return INDEX_HTML


# ============================================================
# 12. START
# ============================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)