# -*- coding: utf-8 -*-

# =========== v2.0 ======================
# 2025.10.20
# This code is the modified version after enabling multi-agent execution.
# The main modification is that batch_run_from_file can read intents from a file
# and generate code sequentially for each intent.
# This code supports multiple rounds of intent parsing and code generation.

from langchain_openai import ChatOpenAI
from typing import Union, Optional, Annotated, Sequence, Dict, Any
from typing import TypedDict, List
from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage, SystemMessage, BaseMessage
from langgraph.graph.message import add_messages
from langchain_core.tools import tool
from langgraph.graph import StateGraph
from typing import TypedDict
from langgraph.graph import START, END
import json
import os
from dotenv import load_dotenv
from typing_extensions import TypedDict
import re


load_dotenv(override=True)
DeepSeek_API_KEY = os.getenv('DEEPSEEK_API_KEY')
model = init_chat_model(model='deepseek-chat', model_provider='deepseek')


class AgentState(TypedDict):
    # ===== General Context =====
    messages: Annotated[Sequence[BaseMessage], add_messages]  # Global conversation history
    user_intent: str  # User modeling intent text
    model_type: str  # Model type, e.g., usecase / requirement / structure / behavior

    # ===== Task Planning Layer =====
    task_cards: List[Dict[str, Any]]  # Current task card list (generated by TaskCardAgent)
    current_task_id: int  # Current task card index being processed
    max_retries: int
    task_card_template: str

    # ===== Code Generation Layer =====
    generated_code: str  # Accumulated generated SysML v2 code
    current_snippet: str  # Code snippet generated from current task card

    # ===== Validation & Repair Layer =====
    validation_result: Optional[Dict[str, Any]]  # Validation results (syntax/structure errors)
    fixed_code: Optional[str]  # Code after repair (if any)
    final_validation_result: Optional[Dict[str, Any]]

    # ===== Management & Monitoring Layer =====
    active_agent: str  # Name of the currently running agent
    progress: Dict[str, float]  # Progress percentage of each stage
    semantic_check: Optional[Dict[str, Any]]  # Final semantic consistency check result


PROMPT_DIR = r"D:\小论文\333\agent2\task_card"

# ========= Model Type Classification =========
def classify_model_type(user_intent: str) -> str:
    text = user_intent.lower()
    if any(k in text for k in ["用例图", "usecase", "use case"]):
        return "usecase"
    elif any(k in text for k in ["需求图", "requirement"]):
        return "requirement"
    elif any(k in text for k in ["结构图", "structure", "block"]):
        return "structure"
    elif any(k in text for k in ["参数图", "约束图", "constraint", "性能", "parametric"]):
        return "parameter"
    elif any(k in text for k in ["状态机图", "活动", "行为", "state", "activity", "行为图", "状态机", "活动图"]):
        return "behavior"

# ========= Template Loading =========
def load_template(model_type: str) -> str:
    template_file = os.path.join(PROMPT_DIR, f"{model_type}.txt")
    if not os.path.exists(template_file):
        raise FileNotFoundError(f"Template file not found: {template_file}")
    with open(template_file, "r", encoding="utf-8") as f:
        return f.read()


GENERAL_INSTANTIATION_PROMPT = """\
You are an intelligent agent (TaskCardAgent) for system-engineering task planning.
Your task is to generate instantiated task cards (JSONL format) based on:
【User Modeling Intent】 and 【General Task Card Template】.
---

【Input Description】

1. task_card_template (template):
   - Contains multiple task-card definitions, each line is a JSON object.
   - Fields include "task_id", "task_name", "sysml_keywords", "objectives", "filled_values".
   - filled_values may contain placeholders such as <PackageName>, <State1>, <Signal1>, etc.

2. user_intent:
   - Describes what system, goal, or process should be modeled.
   - Example: “I want to build a state machine describing the entire process
     from rocket standby, ignition, lift-off, to orbital insertion, including fault detection and reset logic.”
---

【Generation Requirements】

1. You must instantiate each task card using the semantics of user_intent.
2. Replace placeholders in filled_values with concrete values consistent with the intent.
3. Reasonable completion is allowed if the intent does not explicitly mention a field.
4. Do not change field names or structure.
5. Maintain logical ordering and semantic consistency.
6. Output each task card as a JSON object (JSONL format).
7. Output MUST NOT contain explanations, comments, or Markdown.

---

【Output Format】
Each line is a JSON object:
{
  "id": int,
  "O": str,
  "N": str,
  "K": [str],
  "C": str,
  "P": [str],
  "V": {str},
  "D"{
  "depend_on":[str],
  "provides":[str],
  "consumes":[str]
  }
  "objectives": [str],
  "filled_values": {key: value}
}
---

【Begin Now】
Generate instantiated task cards using:

task_card_template:
{{task_card_template}}

user_intent:
{{user_intent}}
"""


def task_card_agent(state: AgentState) -> AgentState:
    """
    TaskCardAgent:
    Generate instantiated task cards based on user intent and model-type templates.
    """

    user_intent = state.get("user_intent", "").strip()
    if not user_intent:
        raise ValueError("❌ Missing user_intent in AgentState.")

    # 1️⃣ Automatically classify model type
    model_type = classify_model_type(user_intent)
    state["model_type"] = model_type
    print(f"🧠 [TaskCardAgent] Model type identified: {model_type}")

    # 2️⃣ Load corresponding template file
    template_filename = f"{model_type}.txt"
    template_path = os.path.join(PROMPT_DIR, template_filename)

    if not os.path.exists(template_path):
        raise FileNotFoundError(f"❌ Template file not found: {template_path}")

    with open(template_path, "r", encoding="utf-8") as f:
        task_card_template = f.read().strip()

    # 3️⃣ Construct full prompt
    full_prompt = (
        GENERAL_INSTANTIATION_PROMPT
        .replace("{{task_card_template}}", task_card_template)
        .replace("{{user_intent}}", user_intent)
    )

    # 4️⃣ Call LLM
    print("🤖 [TaskCardAgent] Generating task cards...")
    response = model.invoke([HumanMessage(content=full_prompt)])
    raw_output = response.content.strip()

    # 5️⃣ Parse JSONL output
    task_cards = []
    for line in raw_output.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            task_cards.append(json.loads(line))
        except json.JSONDecodeError:
            try:
                task_cards.append(eval(line))
            except Exception as e:
                print(f"[Parsing Failed] {line[:100]}... → {e}")

    # 6️⃣ Update AgentState
    state["task_cards"] = task_cards
    state["current_task_id"] = 0
    state["active_agent"] = "TaskCardAgent"
    state["progress"] = {
        "task_card": 1.0,
        "code_gen": 0.0,
        "validation": 0.0,
        "fix": 0.0,
    }

    print(f"✅ [TaskCardAgent] Instantiated {len(task_cards)} task cards.")
    return state


CODEGEN_PROMPT = """\
You are a professional SysML v2 code generation agent.
Your task is to generate new SysML v2 model code based on:
- Provided task card (task_card)
- Existing model code (previous_code)

【Rules】
- previous_code is the partially generated model.
- You MUST extend it instead of rewriting.
- Maintain correct syntax and logical consistency.
- Ensure all element names are unique.

【Reference Example】
A correct and standard SysML v2 example:

example_code:
{example_code}

【Input】
previous_code:
{previous_code}

task_card:
{task_card}

【Generation Requirements】
1. Do not redefine existing elements.
2. Strictly follow SysML v2 syntax.
3. All additions must satisfy objectives of task_card.
4. Do NOT output explanations, comments, or Markdown.
5. Only output complete SysML v2 code.
"""


def code_gen_agent(state: AgentState) -> AgentState:

    task_cards = state.get("task_cards", [])
    current_id = state.get("current_task_id", 0)
    previous_code = state.get("generated_code", "").strip()

    if not task_cards:
        raise ValueError("❌ No task cards available in AgentState.")
    if current_id >= len(task_cards):
        print("✅ All task cards completed.")
        return state

    current_card = task_cards[current_id]
    print(f"🧩 [CodeGenAgent] Current task card: {current_card.get('task_name', 'Unknown Task')}")

    task_card_str = json.dumps(current_card, ensure_ascii=False, indent=2)
    model_type = state["model_type"]

    example_code_filename = f"{model_type}_code.txt"
    code_path = os.path.join(PROMPT_DIR, example_code_filename)

    if not os.path.exists(code_path):
        raise FileNotFoundError(f"❌ Template file not found: {code_path}")

    with open(code_path, "r", encoding="utf-8") as f:
        example_code = f.read().strip()

    full_prompt = CODEGEN_PROMPT.format(
        example_code=example_code,
        previous_code=previous_code if previous_code else "(empty)",
        task_card=task_card_str
    )

    print("🤖 [CodeGenAgent] Generating code...")
    messages = state["messages"] + [HumanMessage(content=full_prompt)]
    response = model.invoke(messages)

    state["messages"].append(HumanMessage(content=full_prompt))
    state["messages"].append(AIMessage(content=response.content))

    new_code = response.content.strip()

    if previous_code and new_code.startswith(previous_code):
        combined_code = new_code
    else:
        combined_code = (
            previous_code + "\n\n" + new_code if previous_code else new_code
        )

    state["generated_code"] = combined_code
    state["current_snippet"] = new_code
    state["current_task_id"] = current_id + 1
    state["active_agent"] = "CodeGenAgent"
    state["progress"]["code_gen"] = round(
        (current_id + 1) / len(task_cards), 2
    )

    print(f"✅ [CodeGenAgent] Completed {current_id + 1}/{len(task_cards)} code layers.")
    return state


def run_interactive_build(user_intent: str):
    """
    Starting from user intent, generate task cards and produce code layer-by-layer.
    """
    print("🚀 Starting SysML v2 auto-modeling demonstration\n")

    state: AgentState = {
        "messages": [],
        "user_intent": user_intent,
        "task_card_template": "",
        "task_cards": [],
        "current_task_id": 0,
        "generated_code": "",
        "current_snippet": "",
        "validation_result": None,
        "fixed_code": None,
        "active_agent": "",
        "progress": {"task_card": 0.0, "code_gen": 0.0},
        "semantic_check": None,
        "model_type": "",
    }

    print("📘 Step 1: Generating task cards...")
    state = task_card_agent(state)
    num_cards = len(state["task_cards"])
    print(f"✅ Generated {num_cards} task cards.\n")

    print("🧩 Step 2: Generating code layer-by-layer...\n")

    while state["current_task_id"] < num_cards:
        state = code_gen_agent(state)
        layer = state["current_task_id"]
        print(f"\n====== Code Layer {layer} ======\n")
        print(state["generated_code"])
        print("==============================\n")

        input("👉 Press Enter for next layer...\n")

    print("🎯 All task cards completed.\n")

    print("================ Final Code ================\n")
    print(state["generated_code"])
    print("=============================================")
    return state


import tempfile
from sysml_validator import validate_sysml_file


def verifier_agent(state: AgentState) -> AgentState:
    """
    VerifierAgent:
    Validate current-layer SysML v2 code for syntax and structure using
    external validate_sysml_file tool.
    """

    current_id = state.get("current_task_id", 0)
    snippet = state.get("current_snippet", "").strip()

    if not snippet:
        print(f"⚠️ [VerifierAgent] Layer {current_id} has no code to validate.")
        state["validation_result"] = {
            "valid": False,
            "errors": [{"severity": "error", "message": f"Layer {current_id} is empty, cannot validate"}],
            "warnings": [],
        }
        return state

    with tempfile.NamedTemporaryFile(delete=False, suffix=".sysml", mode="w", encoding="utf-8") as tmp:
        tmp.write(snippet)
        tmp_path = tmp.name

    print(f"🧪 [VerifierAgent] Validating layer {current_id}...")
    try:
        result = validate_sysml_file(tmp_path, include_warnings=True)
    except Exception as e:
        result = {"valid": False, "errors": [{"severity": "error", "message": str(e)}], "warnings": []}

    try:
        os.remove(tmp_path)
    except OSError:
        pass

    state["validation_result"] = {
        "valid": result.get("valid", False),
        "errors": result.get("errors", []),
        "warnings": result.get("warnings", []),
    }

    if state["validation_result"]["valid"]:
        print(f"✅ [VerifierAgent] Layer {current_id} passed validation.")
    else:
        num_err = len(state["validation_result"].get("errors", []))
        print(f"❌ [VerifierAgent] Layer {current_id} found {num_err} errors.")

    state["active_agent"] = "VerifierAgent"
    state["progress"]["validation"] = round(current_id / max(1, len(state.get("task_cards", []))), 2)

    return state


def fixer_agent(state: AgentState) -> AgentState:
    """
    FixerAgent:
    Repair SysML v2 code based on validation results.
    Apply minimal modifications while keeping original structure consistent.
    """

    current_id = state.get("current_task_id", 0)
    validation = state.get("validation_result", {})
    code_to_fix = state.get("current_snippet", "").strip()

    if not code_to_fix:
        print(f"⚠️ [FixerAgent] Layer {current_id} has no code to repair.")
        return state

    if not validation or validation.get("valid", True):
        print(f"✅ [FixerAgent] Layer {current_id} requires no repair (already valid).")
        state["fixed_code"] = code_to_fix
        return state

    system_prompt = f"""\
You are a SysML v2 model repair expert (FixerAgent).

【Input】
1. Original SysML v2 code
2. Diagnostic results from validator (syntax and structure errors)

【Repair Principles】
1. Apply minimal necessary modifications
2. Preserve existing structure, naming, and hierarchy
3. Fix syntax errors, unclosed brackets, missing semicolons, incorrect transitions, etc.
4. Output the fully repaired SysML v2 code
5. Do NOT output explanations, comments, or Markdown
6. If no errors exist, output the code unchanged

【Original Code】
{code_to_fix}

【Diagnostic Information】
{json.dumps(validation, ensure_ascii=False, indent=2)}

Please output ONLY the repaired SysML v2 code:
"""

    print(f"🩺 [FixerAgent] Repairing layer {current_id}...")
    response = model.invoke([SystemMessage(content=system_prompt)])
    fixed_code = response.content.strip()

    state["fixed_code"] = fixed_code
    state["active_agent"] = "FixerAgent"
    state["progress"]["fix"] = round(current_id / max(1, len(state.get("task_cards", []))), 2)

    print(f"🔧 [FixerAgent] Layer {current_id} repair completed.")
    return state


def orchestrator_agent(state: 'AgentState') -> 'AgentState':
    """
    OrchestratorAgent:
    Controls full pipeline: CodeGenAgent → VerifierAgent → FixerAgent.
    Objective:
      - Generate SysML v2 code layer-by-layer
      - Validate each layer
      - Automatically repair errors
      - Retry until valid or max retries reached
    """

    print("\n🚀 [OrchestratorAgent] Starting code generation & repair workflow...")

    max_retries = state.get("max_retries")
    task_cards = state.get("task_cards", [])
    total_layers = len(task_cards)

    if total_layers == 0:
        print("❌ [OrchestratorAgent] No task cards. Run TaskCardAgent first.")
        return state

    for idx in range(total_layers):
        state["current_task_id"] = idx
        retries = 0
        task_name = task_cards[idx].get("task_name", f"Task {idx+1}")

        print(f"\n[OrchestratorAgent] === Starting Layer {idx+1}/{total_layers}: {task_name} ===")

        state = code_gen_agent(state)
        state = verifier_agent(state)

        while (
            not state.get("validation_result", {}).get("valid", False)
            and retries < max_retries
        ):
            retries += 1
            print(f"[OrchestratorAgent] ❌ Layer {idx+1} invalid, retrying repair ({retries}/{max_retries})")
            state = fixer_agent(state)
            state["current_snippet"] = state.get("fixed_code", "")
            state = verifier_agent(state)

        if state.get("validation_result", {}).get("valid", False):
            print(f"✅ [OrchestratorAgent] Layer {idx+1} validated.")
        else:
            print(f"⚠️ [OrchestratorAgent] Layer {idx+1} exceeded max retries, skipping.")

        print(f"[OrchestratorAgent] === Layer {idx+1} completed ===\n")

    print("\n🎯 [OrchestratorAgent] All layers completed.")
    state["final_validation_result"] = state.get("validation_result", {})
    state["active_agent"] = "OrchestratorAgent"
    return state


def semantic_check_agent_global(state: AgentState) -> AgentState:


    print("\n🧠 [SemanticCheckAgent_Global] Starting global semantic consistency check...")

    task_cards = state.get("task_cards", [])
    final_code = state.get("current_snippet", "").strip()

    if not task_cards or not final_code:
        state["semantic_check"] = {
            "match_score": 0.0,
            "missing_tasks": [],
            "summary": "Missing task cards or final code. Cannot perform semantic check."
        }
        print("⚠️ [SemanticCheckAgent_Global] Missing inputs, skipping.")
        return state

    task_cards_str = json.dumps(task_cards, ensure_ascii=False, indent=2)

    SEMANTIC_PROMPT = f"""
You are a SysML v2 semantic consistency expert.
Compare the task card set with the final generated SysML v2 code
and determine whether all semantic objectives were implemented.

【Task Cards】
{task_cards_str}

【Final SysML v2 Code】
{final_code}

【Analysis Requirements】
1. Check whether the code contains all sysml_keywords declared in task cards.
2. Check whether key elements in filled_values appear in the code.
3. Identify unimplemented or partially implemented task cards.
4. Compute total match score (0–1).
5. Output JSON with the format:

{{
  "match_score": float,
  "missing_tasks": [int],
  "missing_keywords": [str],
  "summary": "Semantic coverage analysis"
}}
"""

    print("🤖 [SemanticCheckAgent_Global] Analyzing...")
    resp = model.invoke([HumanMessage(content=SEMANTIC_PROMPT)])
    raw_output = resp.content.strip()

    try:
        match = re.search(r'\{.*\}', raw_output, re.S)
        if match:
            result = json.loads(match.group())
        else:
            raise ValueError("No JSON found")
    except Exception:
        result = {
            "match_score": 0.0,
            "missing_tasks": [],
            "missing_keywords": [],
            "summary": f"⚠️ Model output not JSON: {resp.content[:200]}..."
        }

    state["semantic_check"] = result
    state["active_agent"] = "SemanticCheckAgent_Global"

    print(f"✅ Consistency Score: {result.get('match_score', 0.0)}")
    print(f"🧩 Missing Task Cards: {result.get('missing_tasks', [])}")
    print(f"🔍 Missing Keywords: {result.get('missing_keywords', [])}")
    print(f"🧾 Summary: {result.get('summary', '')}\n")

    return state


# Build workflow graph
graph = StateGraph(AgentState)

graph.add_node("task_card_agent", task_card_agent)
graph.add_node("orchestrator_agent", orchestrator_agent)
graph.add_node("code_gen_agent", code_gen_agent)
graph.add_node("verifier_agent", verifier_agent)
graph.add_node("fixer_agent", fixer_agent)
graph.add_node("semantic_check_global", semantic_check_agent_global)

graph.add_edge(START, "task_card_agent")
graph.add_edge("task_card_agent", "orchestrator_agent")
graph.add_edge("orchestrator_agent", "semantic_check_global")
graph.add_edge("semantic_check_global", END)

pipeline = graph.compile()


import os
import json
from datetime import datetime
from langchain_core.messages import HumanMessage

def batch_run_from_file(intent_file: str):
    """
    Read multiple modeling intents from specified txt file
    and process each intent sequentially.
    """

    if not os.path.exists(intent_file):
        raise FileNotFoundError(f"❌ Input file not found: {intent_file}")

    with open(intent_file, "r", encoding="utf-8") as f:
        lines = [line.strip() for line in f.readlines() if line.strip()]

    print(f"\n📘 Loaded {len(lines)} modeling intents.\n")

    output_dir = r"D:\333\agent2\output"
    os.makedirs(output_dir, exist_ok=True)

    for idx, intent in enumerate(lines, start=1):
        print(f"\n🚀 [Batch {idx}/{len(lines)}] Processing intent: {intent}\n")

        state: AgentState = {
            "messages": [HumanMessage(content=intent)],
            "user_intent": intent,
            "model_type": "",
            "task_cards": [],
            "task_card_template": "",
            "current_task_id": 0,
            "generated_code": "",
            "current_snippet": "",
            "validation_result": None,
            "fixed_code": None,
            "active_agent": "",
            "progress": {"task_card": 0.0, "code_gen": 0.0, "validation": 0.0, "fix": 0.0},
            "semantic_check": None,
            "max_retries": 3,
            "final_validation_result": None,
        }

        try:
            state = pipeline.invoke(state)
        except Exception as e:
            print(f"❌ [Batch {idx}] Failed to generate: {e}")
            continue

        save_data = {
            "user_intent": state.get("user_intent", ""),
            "task_cards": state.get("task_cards", []),
            "current_snippet": state.get("current_snippet", ""),
            "semantic_check": state.get("semantic_check", {}),
            "final_validation_result": state.get("final_validation_result", []),
        }

        file_name = f"{idx:02d}.json"
        output_path = os.path.join(output_dir, file_name)

        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(save_data, f, ensure_ascii=False, indent=2)

        print(f"✅ [Batch {idx}] Saved to: {output_path}\n")

    print("\n🎯 All modeling intents processed!")
    print(f"📂 Output directory: {output_dir}")


if __name__ == "__main__":
    intent_file = r"D:\小论文\333\agent2\intents.txt"
    batch_run_from_file(intent_file)