"""
Component 5: Self-Refine Feedback Loop
Implements the Self-Refine pattern (arxiv:2303.17651).
"""

import os
import logging
from typing import Dict, List, Optional, Generator, Tuple
from huggingface_hub import InferenceClient

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "Qwen/Qwen2.5-72B-Instruct"
DEFAULT_THRESHOLD = 85
MAX_ITERATIONS = 4


def _get_client(hf_token=None):
    token = hf_token or os.environ.get("HF_TOKEN", "")
    return InferenceClient(api_key=token, timeout=120)


REFINE_SYSTEM_PROMPT = """You are an expert resume optimizer. You have been given:
1. A candidate's generated resume
2. ATS scoring feedback with specific issues

YOUR TASK: Refine the resume to address EVERY piece of feedback while maintaining truthfulness.

RULES:
- ONLY use information already in the resume. Do NOT fabricate.
- Add missing keywords naturally where they genuinely relate to existing skills/experience.
- If a keyword has NO relation to the candidate, do NOT force it.
- You CAN remove irrelevant content to make room for relevant keywords.
- Output the COMPLETE refined resume in plain text. NO markdown.
"""


def _build_refine_prompt(current_cv, ats_report, job_description, iteration):
    feedback_text = "\n".join(f"- {fb}" for fb in ats_report.get("feedback", []))
    kw_details = ats_report.get("breakdown", {}).get("keyword_match", {}).get("details", {})
    missing_kws = kw_details.get("required_missing", [])
    missing_kw_text = ", ".join(missing_kws[:10]) if missing_kws else "None"

    skill_details = ats_report.get("breakdown", {}).get("skill_coverage", {}).get("details", {})
    missing_skills = skill_details.get("missing", [])
    missing_skills_text = ", ".join(missing_skills[:8]) if missing_skills else "None"

    section_details = ats_report.get("breakdown", {}).get("section_completeness", {}).get("details", {})
    missing_sections = section_details.get("required_missing", [])
    missing_sections_text = ", ".join(s.title() for s in missing_sections) if missing_sections else "None"

    return f"""REFINEMENT ITERATION {iteration}/{MAX_ITERATIONS}

CURRENT ATS SCORE: {ats_report.get('total_score', 0)}/100 (Target: {DEFAULT_THRESHOLD}+)

ISSUES TO FIX:
{feedback_text}

MISSING KEYWORDS: {missing_kw_text}
MISSING SKILLS: {missing_skills_text}
MISSING SECTIONS: {missing_sections_text}

CURRENT RESUME:
{current_cv}

JOB DESCRIPTION:
{job_description[:600]}

Rewrite the COMPLETE resume addressing all issues. Output ONLY the improved resume."""


def _safe_stream(client, model, system_prompt, user_prompt, max_tokens=3000):
    """Stream LLM response with safe chunk handling."""
    try:
        stream = client.chat_completion(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            max_tokens=max_tokens,
            temperature=0.3,
            stream=True,
        )
        for chunk in stream:
            try:
                if chunk.choices and len(chunk.choices) > 0:
                    delta = chunk.choices[0].delta
                    if delta and hasattr(delta, 'content') and delta.content:
                        yield delta.content
            except (IndexError, AttributeError):
                continue
    except Exception as e:
        logger.error(f"LLM error in refine: {e}")
        yield f"\n\n⚠️ LLM error: {str(e)}"


def refine_loop(
    initial_cv, parsed_cv, keywords_data, job_description,
    hf_token=None, model=DEFAULT_MODEL, threshold=DEFAULT_THRESHOLD, max_iterations=MAX_ITERATIONS,
):
    """Self-Refine loop. Yields (status_message, current_cv) tuples."""
    from ats_checker import run_ats_check, format_ats_report

    client = _get_client(hf_token)
    current_cv = initial_cv
    best_cv = initial_cv
    best_score = 0

    for iteration in range(1, max_iterations + 1):
        yield (f"🔄 Iteration {iteration}/{max_iterations}: Scoring...", current_cv)

        ats_report = run_ats_check(current_cv, job_description, keywords_data)
        score = ats_report["total_score"]
        report_text = format_ats_report(ats_report)
        logger.info(f"Refine iteration {iteration}: score={score}")

        if score > best_score:
            best_score = score
            best_cv = current_cv

        yield (f"🔄 Iteration {iteration}/{max_iterations}: Score = {score}/100\n\n{report_text}", current_cv)

        if score >= threshold:
            yield (f"✅ Target reached! Score: {score}/100 in {iteration} iteration(s).\n\n{report_text}", current_cv)
            return

        if iteration < max_iterations:
            yield (f"🔄 Iteration {iteration}/{max_iterations}: Score {score} < {threshold}. Refining...", current_cv)

            refine_prompt = _build_refine_prompt(current_cv, ats_report, job_description, iteration)

            refined_cv = ""
            for chunk in _safe_stream(client, model, REFINE_SYSTEM_PROMPT, refine_prompt):
                refined_cv += chunk
                yield (f"🔄 Iteration {iteration}/{max_iterations}: Refining...", refined_cv)

            if refined_cv.strip() and len(refined_cv.strip()) > 100:
                current_cv = refined_cv.strip()
            else:
                yield (f"⚠️ Iteration {iteration}: Empty output. Keeping previous.", current_cv)

    final_report = run_ats_check(best_cv, job_description, keywords_data)
    final_text = format_ats_report(final_report)
    yield (f"🏁 Done. Best score: {best_score}/100\n\n{final_text}", best_cv)