import os
import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor, LogitsProcessorList
from bridge import LatentBridge

class BridgeDecayProcessor(LogitsProcessor):
    def __init__(self, bridge):
        self.bridge = bridge

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        self.bridge.step_generation()
        return scores

def load_project_files(directory, max_chars=20000):
    """
    Load source files from a directory.
    Limits characters to avoid exceeding the context window or VRAM.
    """
    content = ""
    valid_extensions = ('.py', '.ts', '.tsx', '.java', '.js')
    
    for root, dirs, files in os.walk(directory):
        # Ignore useless directories
        if any(ignored in root for ignored in ['node_modules', '.git', '__pycache__', 'dist', 'build']):
            continue
            
        for file in files:
            if file.endswith(valid_extensions):
                filepath = os.path.join(root, file)
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        file_content = f.read()
                        if len(content) + len(file_content) > max_chars:
                            content += f"\n--- {file} (TRUNCATED) ---\n"
                            return content # Limit reached
                        content += f"\n--- {file} ---\n{file_content}\n"
                except Exception:
                    pass
    return content

def main():
    model_name = "Qwen/Qwen3.5-4B"
    checkpoint_path = "bridge_weights.pt"
    

    target_dir = r".." 
    
    print(f"[INFO] Loading Base Model: {model_name} in FP16...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        attn_implementation="sdpa"
    ).to("cuda")
    
    print(f"[INFO] Loading Latent Bridge...")
    bridge = LatentBridge(hidden_dim=model.config.hidden_size, target_layers=[11, 19, 27])
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    bridge.load_state_dict(ckpt["state_dict"] if "state_dict" in ckpt else ckpt)
    device = next(model.parameters()).device
    bridge = bridge.to(device)


    print(f"\n[INFO] Loading source files from: {os.path.abspath(target_dir)}...")
    project_code = load_project_files(target_dir, max_chars=25000)
    print(f"[INFO] Code loaded: ~{len(project_code)} characters.")

    sys_A = "System: You are a Senior Software Architect. Explain the architecture of the provided code clearly and schematically. Describe how the various components interact. Do not think out loud, just produce the final report."
    sys_B = "System: You are an expert code reviewer (Staff Engineer). Deeply study the code, analyze data flows, class responsibilities, and memorize business logic and architectural interactions. Build a complex mental map."
    
    user_prompt = f"User: Analyze the following source code and explain its architecture in detail.\n\nCodebase:\n{project_code}\n\nWhat is the architecture of the system?"


    msgs_A = [{"role": "system", "content": sys_A}, {"role": "user", "content": user_prompt}]
    msgs_B = [{"role": "system", "content": sys_B}, {"role": "user", "content": user_prompt}]
    
    prompt_A = tokenizer.apply_chat_template(msgs_A, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    prompt_B = tokenizer.apply_chat_template(msgs_B, tokenize=False, add_generation_prompt=True, enable_thinking=False)

    inputs_A = tokenizer(prompt_A, return_tensors="pt").to(device)
    inputs_B = tokenizer(prompt_B, return_tensors="pt").to(device)

    print("\n[INFO] === PHASE 1: Background Architecture Study (Agent B) ===")
    bridge.detach()
    bridge.clear_context()
    with torch.no_grad():
        outputs_B = model(**inputs_B, output_hidden_states=True)
    bridge.set_context(outputs_B.hidden_states)
    print("[SUCCESS] Mental map of the code captured in Latent Space!")

    print("\n[INFO] === PHASE 2: Explanation (Agent A guided by Latent Bridge) ===")
    bridge.enable_generation_mode(decay_rate=0.85)
    bridge.attach(model)
    
    processors = LogitsProcessorList([BridgeDecayProcessor(bridge)])

    start_time = time.time()
    with torch.no_grad():
        out = model.generate(
            **inputs_A,
            max_new_tokens=4096,
            do_sample=False,
            repetition_penalty=1.1,
            logits_processor=processors,
            pad_token_id=tokenizer.eos_token_id
        )
    end_time = time.time()
    
    bridge.disable_generation_mode()
    bridge.detach()

    new_tokens = out[0][inputs_A.input_ids.shape[1]:]
    generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    
    generated_text = generated_text.strip()

    print(f"\n[SUCCESS] Generated in {end_time - start_time:.2f} sec.")
    print("\nARCHITECTURAL ANALYSIS (4B with Telepathy):\n" + "="*80)
    print(generated_text)
    print("="*80)

if __name__ == "__main__":
    main()