import json
from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel, Field
from firebase_admin import firestore
from auth import verify_firebase_token
from utils.rate_limiter import with_rate_limit
from utils.groq_client import call_groq_json, HEAVY_MODEL
from utils.pinecone_client import query_book_context
from utils import niche_requests

router = APIRouter()


# ─── Request schemas ──────────────────────────────────────────────────────────

class GenerateFollowupsRequest(BaseModel):
    businessId: str
    section: int          # 1, 2, or 3
    sectionData: dict


class GenerateProfileRequest(BaseModel):
    userId: str
    businessId: str
    force: bool = False   # regenerate outputs even if they already exist (Settings edits)


class ReopenFollowupsRequest(BaseModel):
    businessId: str


class NicheRequestIn(BaseModel):
    text: str = Field(min_length=2, max_length=80)
    sectorId: str = Field(default="", max_length=60)
    subsectorId: str = Field(default="", max_length=60)


# ─── Helpers ──────────────────────────────────────────────────────────────────

def _db():
    return firestore.client()


def _get_bff_data(user_id: str, business_id: str) -> dict:
    """Reads all BFF sections from Firestore for a given business."""
    db = _db()
    biz_ref = (
        db.collection("user_data")
        .document(user_id)
        .collection("businesses")
        .document(business_id)
    )
    sections = {}
    for section_id in [
        "section1", "section1_followups",
        "section2", "section2_followups",
        "section3", "section3_followups",
    ]:
        doc = biz_ref.collection("bff").document(section_id).get()
        if doc.exists:
            sections[section_id] = doc.to_dict()
    return sections


# ─── Business stage ───────────────────────────────────────────────────────────
# Keep ids in sync with data/business-stages.ts on the frontend.

BUSINESS_STAGES = {
    "idea":        "Just an idea — has not started anything yet (Ideation stage)",
    "validation":  "Checking if the idea will work — researching and talking to possible customers (Validation / Planning stage)",
    "prelaunch":   "Getting ready to launch — registering, raising money, setting up (Pre-launch stage)",
    "launched":    "Just started — selling for less than 1 year (Launch / Startup stage)",
    "early":       "Running for 1–2 years (Early stage)",
    "growth":      "Running for 3–5 years (Growth stage)",
    "established": "Running for 6–10 years (Established / Expansion stage)",
    "mature":      "Running for over 10 years (Mature stage)",
    "restarting":  "Restarting or changing direction — has run a business before and is starting again or pivoting",
}

PRE_LAUNCH_STAGES = {"idea", "validation", "prelaunch"}


def _get_business_stage(user_id: str, business_id: str) -> str:
    """Reads the stage saved on the WHAT step. Returns '' if not set."""
    snap = (
        _db().collection("user_data").document(user_id)
        .collection("businesses").document(business_id)
        .collection("bff").document("section1").get()
    )
    return (snap.to_dict() or {}).get("businessStage", "") if snap.exists else ""


def _stage_guidance(stage: str) -> str:
    """Prompt instructions that keep the AI's tense and assumptions matched to the stage."""
    if stage not in BUSINESS_STAGES:
        return ""
    label = BUSINESS_STAGES[stage]
    if stage in PRE_LAUNCH_STAGES:
        rule = (
            "This business has NOT launched yet. Never assume it has customers, sales, staff, "
            "revenue or a track record. Phrase everything around plans, expectations and "
            "assumptions (e.g. 'Who do you expect to buy this?' not 'Who buys from you?')."
        )
    elif stage == "restarting":
        rule = (
            "The entrepreneur has run a business before and is starting again or changing direction. "
            "Treat the new direction as not yet proven: phrase it around plans, but draw on their "
            "past experience (what worked, what didn't, what they are changing and why)."
        )
    else:
        rule = (
            "This business is already operating. Ask about real experience so far — actual "
            "customers, what sells, what has worked and what hasn't — at a depth that suits "
            "how long it has been running."
        )
    return f"\n\nBUSINESS STAGE: {label}\n{rule}"


def _call_groq_json(system_prompt: str, user_prompt: str, fixture: str | None = None) -> dict:
    """
    Thin wrapper — delegates to the key-rotating Groq client utility.

    max_tokens is generous (3000) because HEAVY_MODEL is a reasoning model
    (openai/gpt-oss-120b) that spends part of its budget on internal
    reasoning before writing the final JSON — a tight budget here truncates
    the JSON output and fails to parse.
    """
    return call_groq_json(system_prompt, user_prompt, model=HEAVY_MODEL, max_tokens=3000, fixture=fixture)


# ─── Route: Generate follow-up questions ─────────────────────────────────────

@router.post("/generate-followups")
def generate_followups(
    body: GenerateFollowupsRequest,
    user: dict = Depends(with_rate_limit("generation")),
):
    """
    After each BFF section is submitted, generates 3–5 personalized
    follow-up questions specific to the user's sector and answers.
    Questions use plain English — no business jargon.
    """
    if body.section not in (1, 2, 3):
        raise HTTPException(status_code=400, detail="section must be 1, 2, or 3.")

    section_names = {1: "WHAT (identity)", 2: "WHY (purpose)", 3: "HOW (operations)"}

    # Fetch relevant BFF philosophy passages from "Built on Rock"
    bff_query = f"Business Footing Framework section {body.section} {json.dumps(body.sectionData)[:200]}"
    book_ctx   = query_book_context(bff_query, top_k=2)

    # Stage is read server-side so WHY/HOW follow-ups get it too, not just WHAT
    stage = body.sectionData.get("businessStage") if body.section == 1 else None
    stage = stage or _get_business_stage(user["uid"], body.businessId)

    system_prompt = (
        "You are a business consultant helping a Nigerian entrepreneur understand their business better. "
        "Generate follow-up questions based on what they have already shared. "
        "Questions MUST use simple, everyday English — no business jargon. "
        "Think: how would a friendly, experienced mentor ask? "
        "Return ONLY valid JSON."
        + _stage_guidance(stage)
        + (f"\n\n{book_ctx}" if book_ctx else "")
    )

    user_prompt = f"""
The entrepreneur just completed Section {body.section} ({section_names[body.section]}) of their Business Footing Framework.

Their answers:
{json.dumps(body.sectionData, indent=2)}

Generate 3-5 follow-up questions to deepen understanding of their business.
Make each question specific to what they shared — do not ask generic questions.

Return JSON in this exact format:
{{
  "questions": [
    {{
      "id": "q1",
      "question": "...",
      "type": "text",
      "placeholder": "..."
    }}
  ]
}}

Question types:
- "text"        → free-text answer (use for open questions)
- "select"      → single choice, add "options": ["Option A", "Option B", ...]
- "multiselect" → multiple choice, add "options": [...]
- "boolean"     → yes/no question
"""

    try:
        data = _call_groq_json(system_prompt, user_prompt, fixture="followups")
        return data
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Follow-up generation failed: {str(e)}")


# ─── Route: Generate full business profile ────────────────────────────────────

@router.post("/generate-profile")
def generate_profile(
    body: GenerateProfileRequest,
    user: dict = Depends(with_rate_limit("generation")),
):
    """
    Reads all completed BFF sections from Firestore and generates all 4
    profile outputs in parallel (Business Identity, Customer Profile,
    Market Data, Business Model). Saves results back to Firestore and
    marks bffCompleted = True on the business document.
    """
    # Only the document owner can trigger their own profile generation
    if user.get("uid") != body.userId:
        raise HTTPException(status_code=403, detail="Forbidden.")

    bff = _get_bff_data(body.userId, body.businessId)
    if not bff:
        raise HTTPException(status_code=404, detail="BFF data not found.")

    bff_summary = json.dumps(bff, indent=2)

    # ── Retrieve relevant passages from "Built on Rock" via Pinecone ─────────
    sector = bff.get("section1", {}).get("sectorId", "")
    s1     = bff.get("section1", {})
    # nicheId "other" means the person described their own niche in customNiche.
    niche  = s1.get("customNiche") if s1.get("nicheId") == "other" else s1.get("nicheId", "")
    book_ctx = query_book_context(
        f"Business Footing Framework WHAT WHY HOW {sector} {niche} Nigerian entrepreneur",
        top_k=3,
    )

    # ── System prompt used for all 4 calls ───────────────────────────────────
    base_system = (
        "You are an expert business consultant specializing in Nigerian MSMEs, "
        "trained in the Business Footing Framework (BFF) from 'Built on Rock' "
        "by Iyanuoluwa Joseph Akinnusi. "
        "Analyze the BFF data and produce professional, structured profile components. "
        "Write in plain English that ANY entrepreneur can understand. "
        "Ground all analysis in the Nigerian market context. "
        "Return ONLY valid JSON."
        + _stage_guidance(bff.get("section1", {}).get("businessStage", ""))
        + (f"\n\n{book_ctx}" if book_ctx else "")
    )

    # ── Prompt builders ───────────────────────────────────────────────────────
    def prompt_identity():
        return f"""
BFF Data:
{bff_summary}

Generate the Business Identity profile. Return JSON:
{{
  "businessName": "string",
  "sector": "string",
  "subsector": "string",
  "niche": "string",
  "summary": "2-3 sentence plain-English description of what this business does",
  "valueProposition": "one clear sentence: what the business promises customers",
  "brandVoice": "2-3 adjectives describing how the brand should communicate (e.g. warm, direct, trustworthy)",
  "tagline": "a short, memorable tagline for this business"
}}"""

    def prompt_customer():
        return f"""
BFF Data:
{bff_summary}

Generate the Customer Profile. Return JSON:
{{
  "demographics": {{
    "ageRange": "primary age range",
    "gender": "primary gender",
    "location": "primary location(s)",
    "incomeLevel": "low / middle / high / mixed",
    "educationLevel": "brief description",
    "employmentStatus": "brief description"
  }},
  "psychographics": {{
    "values": ["list of 3-5 values the customer likely holds"],
    "lifestyle": "brief lifestyle description",
    "painPoints": ["3-5 main pain points this business solves for them"],
    "motivations": ["3-5 reasons they would buy from this business"]
  }},
  "buyerPersona": "A 3-4 sentence narrative describing the typical customer as a named person (e.g. Meet Amaka...). Make it specific and relatable for Nigerian context."
}}"""

    def prompt_market():
        return f"""
BFF Data:
{bff_summary}

Generate the Market Data analysis. Return JSON:
{{
  "tam": "Total Addressable Market — brief estimate for Nigeria",
  "sam": "Serviceable Addressable Market — realistic target market segment",
  "som": "Serviceable Obtainable Market — what this business can realistically capture",
  "swot": {{
    "strengths": ["3-4 key strengths based on BFF data"],
    "weaknesses": ["3-4 potential weaknesses to address"],
    "opportunities": ["3-4 market opportunities in Nigeria right now"],
    "threats": ["3-4 threats or risks to watch out for"]
  }},
  "pestle": {{
    "political": "brief note on political factors relevant to this sector in Nigeria",
    "economic": "brief note on economic factors",
    "social": "brief note on social/cultural factors",
    "technological": "brief note on tech trends",
    "legal": "brief note on legal/regulatory considerations",
    "environmental": "brief note if applicable"
  }},
  "competitiveLandscape": "2-3 sentences describing who else competes in this space in Nigeria and how this business can differentiate"
}}"""

    def prompt_bmc():
        return f"""
BFF Data:
{bff_summary}

Generate the Business Model Canvas. Return JSON with exactly these 9 keys,
each containing a list of 2-4 concise bullet points:
{{
  "keyPartners": [...],
  "keyActivities": [...],
  "keyResources": [...],
  "valuePropositions": [...],
  "customerRelationships": [...],
  "channels": [...],
  "customerSegments": [...],
  "costStructure": [...],
  "revenueStreams": [...]
}}"""

    # ── Build the call plan — skip outputs already in Firestore ─────────────
    # On retry after a partial failure, this avoids regenerating completed
    # outputs and wasting tokens on the Groq free-tier TPM limit.
    import time
    db_check = _db()
    biz_ref_check = (
        db_check.collection("user_data")
        .document(body.userId)
        .collection("businesses")
        .document(body.businessId)
    )

    all_outputs = {
        "businessIdentity": prompt_identity,
        "customerProfile":  prompt_customer,
        "marketData":       prompt_market,
        "businessModel":    prompt_bmc,
    }

    # profileStale is set by /reopen-followups: the old profile was built without
    # follow-up answers, so it must be rebuilt even if the outputs already exist.
    # Read server-side so it survives the user leaving and resuming later.
    biz_snap = biz_ref_check.get()
    force = body.force or bool((biz_snap.to_dict() or {}).get("profileStale"))

    call_plan = []
    skipped_keys: set = set()
    for key, prompt_fn in all_outputs.items():
        existing = None if force else biz_ref_check.collection("profile_output").document(key).get()
        if existing is not None and existing.exists:
            print(f"[bff] Skipping {key} — already exists in Firestore")
            skipped_keys.add(key)
        else:
            call_plan.append((key, prompt_fn))

    results: dict = {}
    errors: list = []

    # Sequential calls with 15s gap to stay within the free-tier TPM limit.
    # max_tokens=3000 per call leaves headroom for HEAVY_MODEL's internal
    # reasoning tokens on top of the JSON output (see _call_groq_json).
    # 15s gap × 200 tokens/sec refill recovers ~3,000 tokens between calls —
    # if that's not enough, the 5-key rotation pool absorbs the rest via 429s.
    for i, (key, prompt_fn) in enumerate(call_plan):
        if i > 0:
            time.sleep(15)  # 15s gap: ~3,000 tokens recovered at 200 tok/s
        try:
            results[key] = _call_groq_json(base_system, prompt_fn(), fixture=key)
        except Exception as e:
            errors.append(f"{key}: {str(e)}")
            print(f"[bff] Failed to generate {key}: {e}")

    # Fail only if we had work to do and ALL of it failed
    if call_plan and len(errors) == len(call_plan):
        raise HTTPException(
            status_code=500,
            detail=f"Profile generation failed: {'; '.join(errors)}",
        )

    # ── Save to Firestore ─────────────────────────────────────────────────────
    db = _db()
    biz_ref = (
        db.collection("user_data")
        .document(body.userId)
        .collection("businesses")
        .document(body.businessId)
    )

    from google.cloud.firestore_v1 import SERVER_TIMESTAMP
    for key, data in results.items():
        biz_ref.collection("profile_output").document(key).set(
            {**data, "generatedAt": SERVER_TIMESTAMP}
        )

    # Only mark BFF as complete once EVERY output actually exists — either
    # generated just now or already present from a prior partial run. Marking
    # this true when an output is still missing is what previously let a
    # business look "complete" with an empty/partial profile (see
    # project_multibusiness_gate memory — the manual-Console-flip incident
    # had the same shape of bug).
    # On a forced run, outputs that failed to regenerate still exist from before,
    # so count them as present rather than blocking completion.
    if force:
        skipped_keys = {
            k for k in all_outputs
            if k not in results and biz_ref.collection("profile_output").document(k).get().exists
        }
    present_keys = skipped_keys | set(results.keys())
    all_present  = present_keys == set(all_outputs.keys())
    if all_present:
        biz_ref.update({"bffCompleted": True, "profileStale": False})
    else:
        missing = set(all_outputs.keys()) - present_keys
        print(f"[bff] NOT marking bffCompleted — missing outputs: {missing}")

    return {
        "success": all_present,
        "outputs": list(results.keys()),
        "missing": sorted(set(all_outputs.keys()) - present_keys),
    }


# ─── Route: Reopen onboarding for missing follow-ups ─────────────────────────

def _missing_followup_sections(user_id: str, business_id: str) -> list[int]:
    """Sections (1-3) whose follow-up questions were never generated/saved."""
    bff_ref = (
        _db().collection("user_data").document(user_id)
        .collection("businesses").document(business_id)
        .collection("bff")
    )
    missing = []
    for section in (1, 2, 3):
        snap = bff_ref.document(f"section{section}_followups").get()
        if not snap.exists or not (snap.to_dict() or {}).get("questions"):
            missing.append(section)
    return missing


@router.post("/reopen-followups")
def reopen_followups(
    body: ReopenFollowupsRequest,
    user: dict = Depends(verify_firebase_token),
):
    """
    A profile was generated while some follow-up steps failed or were skipped.
    Marks the business incomplete (bffCompleted is backend-only per Firestore
    rules) and flags the profile stale so the next generate-profile rebuilds it
    with the follow-up answers included.
    """
    biz_ref = (
        _db().collection("user_data").document(user["uid"])
        .collection("businesses").document(body.businessId)
    )
    if not biz_ref.get().exists:
        raise HTTPException(status_code=404, detail="Business not found.")

    missing = _missing_followup_sections(user["uid"], body.businessId)
    if missing:
        biz_ref.update({"bffCompleted": False, "profileStale": True})
    return {"missing": missing}


# ─── Route: Niche not in our list ────────────────────────────────────────────

@router.post("/niche-request")
def niche_request(body: NicheRequestIn, user: dict = Depends(verify_firebase_token)):
    """Records a niche the person couldn't find, for admins to review and add."""
    try:
        niche_requests.record(user["uid"], body.text, body.sectorId, body.subsectorId)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return {"success": True}