main.py

import os
from dotenv import load_dotenv
# Load environment variables explicitly from current directory
dotenv_path = os.path.join(os.path.dirname(__file__), '.env')
load_dotenv(dotenv_path)

print(f"DEBUG: GEMINI_API_KEY loaded: {bool(os.getenv('GEMINI_API_KEY'))}")
print(f"DEBUG: GROQ_API_KEY loaded: {bool(os.getenv('GROQ_API_KEY'))}")

import re
import json
import nltk
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.corpus import stopwords
from nltk.probability import FreqDist
from cleantext import clean
import PyPDF2
from textblob import TextBlob
import groq
import traceback
import unidecode
import contractions
from sklearn.feature_extraction.text import TfidfVectorizer
from typing import Optional, Tuple

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

def ensure_nltk_resources():
    import warnings
    import ssl
    try:
        nltk.data.find('tokenizers/punkt')
        nltk.data.find('corpora/stopwords')
    except LookupError:
        try:
            _create_unverified_https_context = ssl._create_unverified_context
        except AttributeError:
            pass
        else:
            ssl._create_default_https_context = _create_unverified_https_context
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            nltk.download('stopwords', quiet=True)
            nltk.download('wordnet', quiet=True)
            nltk.download('words', quiet=True)
            nltk.download('punkt', quiet=True)
            nltk.download('punkt_tab', quiet=True)

ensure_nltk_resources()

# Initialize stopwords
stop_words = set(stopwords.words('english'))
stop_words.update({'ask', 'much', 'thank', 'etc.', 'e', 'We', 'In', 'ed', 'pa', 'This', 'also', 'A', 'fu', 'To', '5', 'ing', 'er', '2'})

_stopwords_list = [w for w in stopwords.words('english') if len(w) >= 3]
_stopwords_pattern = re.compile(r'\b(' + r'|'.join(_stopwords_list) + r')\b\s*')

# Initialize LLM Client logic (from app.py)
class FailoverLLMClient:
    def __init__(self):
        self.groq_api_key = os.getenv("GROQ_API_KEY")
        self.groq_client = groq.Groq(api_key=self.groq_api_key) if self.groq_api_key else None
        self.gemini_api_key = os.getenv("GEMINI_API_KEY")
        print(f"DEBUG: FailoverLLMClient init - Gemini Key: {bool(self.gemini_api_key)}")
        self.gemini_client = None
        if self.gemini_api_key:
            try:
                from google import genai
                self.gemini_client = genai.Client(api_key=self.gemini_api_key)
            except ImportError:
                print("Warning: google-genai not installed.")
            except Exception as e:
                self.gemini_init_error = str(e)
                print(f"Warning: Failed to initialize Gemini client: {e}")
        self._current_provider = None

    def _get_groq_primary(self):
        if self.groq_client: return self.groq_client, "meta-llama/llama-4-scout-17b-16e-instruct"
        return None, ""

    def _get_groq_backup(self):
        if self.groq_client: return self.groq_client, "llama-3.1-8b-instant"
        return None, ""

    def _call_gemini(self, messages: list, temperature: float, max_tokens: int):
        if not self.gemini_client: return None
        from google.genai import types
        prompt = ""
        system_instruction = "You are an expert political science analyst."
        for msg in messages:
            if msg["role"] == "system": system_instruction = msg["content"]
            elif msg["role"] == "user": prompt = msg["content"]
        response = self.gemini_client.models.generate_content(
            model='gemini-2.5-flash',
            contents=prompt,
            config=types.GenerateContentConfig(
                system_instruction=system_instruction,
                temperature=temperature,
                max_output_tokens=max_tokens
            )
        )
        return response.text

    def _call_groq_with_retry(self, model, messages, temperature, max_tokens, retries=2):
        import time
        for attempt in range(retries):
            try:
                params = {
                    "model": model,
                    "messages": messages,
                }
                if "gpt-oss" in model or "llama-4-scout" in model:
                    params["max_completion_tokens"] = 8192
                    if "gpt-oss" in model:
                        params["reasoning_effort"] = "medium"
                    else:
                        params["temperature"] = temperature
                else:
                    params["temperature"] = temperature
                    params["max_tokens"] = max_tokens
                response = self.groq_client.chat.completions.create(**params)
                return response.choices[0].message.content
            except Exception as e:
                error_str = str(e)
                if '429' in error_str and attempt < retries - 1:
                    wait_match = re.search(r'try again in ([\d.]+)s', error_str)
                    wait_time = float(wait_match.group(1)) + 1.0 if wait_match else 30.0
                    wait_time = min(wait_time, 60.0)
                    time.sleep(wait_time)
                    continue
                raise
        return None

    def chat_completion(self, messages, temperature=0.3, max_tokens=1500, operation="operation"):
        providers_tried = []
        if self.groq_client:
            try:
                _, model = self._get_groq_primary()
                if model:
                    result = self._call_groq_with_retry(model, messages, temperature, max_tokens)
                    if result:
                        self._current_provider = f"Groq ({model})"
                        return clean_llm_output(result), self._current_provider
            except Exception as e:
                providers_tried.append(f"Groq ({model}) - FAILED: {str(e)[:150]}")

        # Try Gemini
        if self.gemini_client:
            try:
                result = self._call_gemini(messages, temperature, max_tokens)
                if result:
                    self._current_provider = "Gemini (2.5-flash)"
                    return clean_llm_output(result), self._current_provider
            except Exception as e:
                error_msg = str(e)
                if '429' in error_msg:
                    providers_tried.append(f"Gemini - FAILED: Quota Exhausted (429)")
                else:
                    providers_tried.append(f"Gemini - FAILED: {error_msg[:150]}")
        else:
            reason = getattr(self, 'gemini_init_error', 'Missing Key or Library')
            providers_tried.append(f"Gemini - SKIPPED: {reason}")

        if self.groq_client:
            try:
                _, model = self._get_groq_backup()
                if model:
                    result = self._call_groq_with_retry(model, messages, temperature, max_tokens)
                    if result:
                        self._current_provider = f"Groq ({model}) [BACKUP]"
                        return clean_llm_output(result), self._current_provider
            except Exception as e:
                providers_tried.append(f"Groq ({model}) [BACKUP] - FAILED: {str(e)[:150]}")

        self._current_provider = None
        failure_log = " | ".join(providers_tried)
        
        # Final safety: If Gemini was skipped, make sure the user knows why
        if "Gemini" not in failure_log:
            reason = getattr(self, 'gemini_init_error', 'Not Initialized')
            failure_log += f" | Gemini - SKIPPED: {reason}"
            
        return None, failure_log

    def is_available(self):
        status = []
        if self.groq_client: status.append("Groq: Ready")
        else: status.append("Groq: Missing Key")
        if self.gemini_client: status.append("Gemini: Ready")
        else: 
            err = getattr(self, 'gemini_init_error', 'Missing Key/Library')
            status.append(f"Gemini: {err}")
        return self.groq_client is not None or self.gemini_client is not None, " | ".join(status)

def clean_llm_output(text):
    if not text: return text
    # Strip <think> tags even if they are not closed (e.g. if the response was truncated)
    text = re.sub(r'<think>.*?(?:</think>|$)', '', text, flags=re.DOTALL)
    text = re.sub(r'<br\s*/?>', '\n', text)
    return text.strip()

llm_client = FailoverLLMClient()

def _fix_pdf_text(text):
    # Only remove null bytes - don't try to rejoin words as it corrupts valid words
    text = text.replace('\x00', '')
    # Fix hyphenated line breaks (e.g. "govern-\nment" -> "government")
    text = re.sub(r'(\w)-\n(\w)', r'\1\2', text)
    return text

def parse_pdf(file_bytes):
    from io import BytesIO
    try:
        text = ""
        pdf_file = BytesIO(file_bytes)
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        for page_num in range(len(pdf_reader.pages)):
            page = pdf_reader.pages[page_num]
            text += page.extract_text() + "\n"
        text = _fix_pdf_text(text)
        return clean(text)
    except Exception as e:
        print(f"Error parsing PDF: {e}")
        return None

def clean_text(text):
    text = text.encode("ascii", errors="ignore").decode("ascii")
    text = unidecode.unidecode(text)
    text = contractions.fix(text)
    text = re.sub(r"\n", " ", text)
    text = re.sub(r"\t", " ", text)
    text = re.sub(r"/ ", " ", text)
    text = text.strip()
    text = re.sub(" +", " ", text).strip()
    text = [word for word in text.split() if word not in stop_words]
    return ' '.join(text)

def preprocess(textParty):
    text1Party = re.sub('[^A-Za-z0-9]+', ' ', textParty.lower())
    text2Party = _stopwords_pattern.sub('', text1Party)
    text2Party = ' '.join(w for w in text2Party.split() if len(w) >= 3)
    return text2Party

def generate_summary(text):
    available, status = llm_client.is_available()
    if not available: return f"Summarization is not available. (Status: {status})"
    
    # Smart sampling for large documents: Take chunks from start, middle, and end
    # to ensure the summary covers the entire manifesto.
    full_len = len(text)
    if full_len > 18000:
        chunk_size = 6000
        start_chunk = text[:chunk_size]
        mid_point = full_len // 2
        mid_chunk = text[mid_point - (chunk_size // 2) : mid_point + (chunk_size // 2)]
        end_chunk = text[-chunk_size:]
        text_to_analyze = f"[BEGINNING]\n{start_chunk}\n\n[MIDDLE]\n{mid_chunk}\n\n[END]\n{end_chunk}"
    else:
        text_to_analyze = text

    messages = [
        {"role": "system", "content": "You are an expert political science analyst. Provide a comprehensive, objective, and structured summary of the political manifesto. Use Markdown with headings, bullet points, and bold text for key terms. Do not stop abruptly; ensure your response is complete."},
        {"role": "user", "content": f"Summarize this political manifesto, focusing on main policy areas, promises, and core themes. Aim for 300-500 words.\n\nManifesto Text:\n{text_to_analyze}"}
    ]
    content, provider = llm_client.chat_completion(messages=messages, temperature=0.3, max_tokens=2500)
    return content if content else f"Error generating summary: {provider}"

def get_contextual_search_result(target_word, tar_passage, max_context_length=15000):
    if not target_word or target_word.strip() == "": return "Please enter a search term."
    available, status = llm_client.is_available()
    if not available: return f"Contextual search requires API keys. (Status: {status})"
    
    # Smart context extraction: find up to top 5 occurrences to keep context focused
    search_term_lower = target_word.lower()
    text_lower = tar_passage.lower()
    if search_term_lower in text_lower:
        chunks = []
        start = 0
        window = 1000
        count = 0
        while count < 5:
            idx = text_lower.find(search_term_lower, start)
            if idx == -1: break
            window_start = max(0, idx - window)
            window_end = min(len(tar_passage), idx + window + len(target_word))
            chunks.append(tar_passage[window_start:window_end].strip())
            start = window_end
            count += 1
        
        combined = "\n\n...[relevant section]...\n\n".join(chunks)
        tar_passage_truncated = combined[:max_context_length] if len(combined) > max_context_length else combined
    else:
        tar_passage_truncated = tar_passage[:max_context_length]

    prompt = f"""You are an expert political analyst. You have been given relevant sections of a political manifesto and a specific search term. 
Your task is to provide a very concise summary (exactly 1 to 2 paragraphs) of all information related to the search term.
Focus on:
1. Specific policies, promises, or statements related to the term.
2. Any key details or commitments mentioned.

Search Term: {target_word}
Manifesto Text Sections:
{tar_passage_truncated}

Response (Concise, 1-2 paragraphs):"""
    messages = [
        {"role": "system", "content": "You are an expert political analyst. Provide a concise 1-2 paragraph summary with Markdown formatting."},
        {"role": "user", "content": prompt}
    ]
    result, provider = llm_client.chat_completion(messages=messages, temperature=0.2, max_tokens=1000)
    return result if result else f"Error: {provider}"

def get_f_distance(text2Party):
    if not text2Party or not text2Party.strip(): return {"no_content": 1.0}
    word_tokens_party = word_tokenize(text2Party)
    # Filter: min 4 chars, alpha only, no concatenated junk (no uppercase run-ons)
    word_tokens_party = [
        w for w in word_tokens_party 
        if len(w) >= 4 and w.isalpha() and w == w.lower()
    ]
    if not word_tokens_party: return {"no_tokens": 1.0}
    fdistance = FreqDist(word_tokens_party).most_common(15)
    mem = {x[0]: x[1] for x in fdistance}
    try:
        sentences = sent_tokenize(text2Party)
        if not sentences: return normalize(mem) if mem else {"no_data": 1.0}
        vectorizer = TfidfVectorizer(max_features=15, stop_words='english', token_pattern=r'(?u)\b[A-Za-z]{3,}\b')
        tfidf_matrix = vectorizer.fit_transform(sentences)
        feature_names = vectorizer.get_feature_names_out()
        tfidf_scores = {}
        for i, word in enumerate(feature_names):
            scores = [tfidf_matrix[j, i] for j in range(tfidf_matrix.shape[0]) if i < tfidf_matrix.shape[1]]
            if scores: tfidf_scores[word] = sum(scores) / len(scores)
        combined_scores = {}
        all_words = set(list(mem.keys()) + list(tfidf_scores.keys()))
        max_freq = max(mem.values()) if mem else 1
        max_tfidf = max(tfidf_scores.values()) if tfidf_scores else 1
        for word in all_words:
            combined_scores[word] = (mem.get(word, 0) / max_freq * 0.3) + (tfidf_scores.get(word, 0) / max_tfidf * 0.7)
        top_words = dict(sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)[:10])
        return normalize(top_words) if top_words else (normalize(mem) if mem else {"no_data": 1.0})
    except Exception as e:
        print(f"Error in fDistance: {e}")
        return normalize(mem) if mem else {"error": 1.0}

def normalize(d, target=1.0):
    raw = sum(d.values())
    factor = target / raw if raw != 0 else 0
    return {key: value * factor for key, value in d.items()}

def get_dispersion_data(textParty, top_words):
    word_tokens_party = word_tokenize(textParty.lower())
    if not word_tokens_party or not top_words: return {}
    dispersion_data = {}
    for word in top_words:
        target = word.lower()
        offsets = [j for j, token in enumerate(word_tokens_party) if token == target]
        dispersion_data[word] = offsets
    return dispersion_data

# --- FastAPI App ---

app = FastAPI(title="Manifesto Explainer API")

print("DEBUG: Registering /test_env route")
@app.get("/test_env")
async def test_env():
    return {
        "groq_key_loaded": bool(os.getenv("GROQ_API_KEY")),
        "gemini_key_loaded": bool(os.getenv("GEMINI_API_KEY")),
        "groq_key_prefix": os.getenv("GROQ_API_KEY")[:5] if os.getenv("GROQ_API_KEY") else None,
        "gemini_key_prefix": os.getenv("GEMINI_API_KEY")[:5] if os.getenv("GEMINI_API_KEY") else None
    }


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/analyze")
async def analyze_manifesto(
    file: Optional[UploadFile] = File(None), 
    example_file: Optional[str] = Form(None),
    search_term: str = Form("government")
):
    try:
        if example_file:
            valid_examples = ["Example/AAP_Manifesto_2019.pdf", "Example/Bjp_Manifesto_2019.pdf", "Example/Congress_Manifesto_2019.pdf"]
            # Convert paths to use backslashes or forward slashes appropriately just in case, but keep simple check
            is_valid = any(ex in example_file.replace('\\', '/') for ex in valid_examples)
            if not is_valid:
                raise HTTPException(status_code=400, detail="Invalid example file.")
            
            # Check if file actually exists
            if not os.path.exists(example_file):
                raise HTTPException(status_code=400, detail=f"Example file not found: {example_file}")
                
            with open(example_file, "rb") as f:
                file_bytes = f.read()
        elif file:
            file_bytes = await file.read()
        else:
            raise HTTPException(status_code=400, detail="No file or example provided.")
            
        raw_party = parse_pdf(file_bytes)
        
        if not raw_party:
            raise HTTPException(status_code=400, detail="Failed to parse PDF.")
        
        text_party = clean_text(raw_party)
        text_party_processed = preprocess(text_party)
        
        MAX_VIS_CHARS = 30000
        text_for_vis = text_party_processed[:MAX_VIS_CHARS] if len(text_party_processed) > MAX_VIS_CHARS else text_party_processed
        
        # 1. Summary
        summary = generate_summary(raw_party)
        
        # 2. Contextual Search
        search_res = get_contextual_search_result(search_term, raw_party)
        
        # 3. Frequency / Topics
        fdist_party = get_f_distance(text_for_vis)
        
        # 4. Sentiment & Subjectivity
        if not text_for_vis.strip():
            polarity_val = 0.0
            subjectivity_val = 0.0
        else:
            polarity_val = TextBlob(text_for_vis).sentiment.polarity
            subjectivity_val = TextBlob(text_for_vis).sentiment.subjectivity
            
        # 5. Dispersion Data
        top_5_words = list(fdist_party.keys())[:5]
        dispersion_data = get_dispersion_data(text_for_vis, top_5_words)
        
        # Word cloud raw frequencies (top 150 words)
        word_tokens_party = [w for w in word_tokenize(text_for_vis) if len(w) >= 3 and w.isalpha()]
        word_cloud_freq = dict(FreqDist(word_tokens_party).most_common(150))
        
        return JSONResponse(content={
            "summary": summary,
            "search_result": search_res,
            "topics": fdist_party,
            "sentiment": {
                "polarity": polarity_val,
                "subjectivity": subjectivity_val
            },
            "dispersion": dispersion_data,
            "word_cloud_freq": word_cloud_freq,
            "total_tokens": len(word_tokenize(text_party.lower()))
        })

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

print("DEBUG: Registering /test_env route")
@app.get("/test_env")
async def test_env():
    return {
        "groq_key_loaded": bool(os.getenv("GROQ_API_KEY")),
        "gemini_key_loaded": bool(os.getenv("GEMINI_API_KEY")),
        "groq_key_prefix": os.getenv("GROQ_API_KEY")[:5] if os.getenv("GROQ_API_KEY") else None,
        "gemini_key_prefix": os.getenv("GEMINI_API_KEY")[:5] if os.getenv("GEMINI_API_KEY") else None
    }

# --- Serve Static Frontend Files ---(no-cache so edits always load fresh) ---

NO_CACHE_HEADERS = {
    "Cache-Control": "no-store, no-cache, must-revalidate",
    "Pragma": "no-cache",
}

@app.get("/")
async def serve_index():
    return FileResponse("index.html", headers=NO_CACHE_HEADERS)

@app.get("/style.css")
async def serve_css():
    return FileResponse("style.css", media_type="text/css", headers=NO_CACHE_HEADERS)

@app.get("/script.js")
async def serve_js():
    return FileResponse("script.js", media_type="application/javascript", headers=NO_CACHE_HEADERS)

# Serve any other static assets (e.g. images)
app.mount("/static", StaticFiles(directory="."), name="static")

if __name__ == "__main__":
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)